FiTv2: Масштабируемый и улучшенный гибкий видовой трансформер для моделирования диффузии

FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model

October 17, 2024

Авторы: ZiDong Wang, Zeyu Lu, Di Huang, Cai Zhou, Wanli Ouyang, and Lei Bai

cs.AI

Аннотация

Природа бесконечно свободна от разрешения. В контексте этой реальности существующие модели диффузии, такие как Диффузионные Трансформеры, часто сталкиваются с проблемами при обработке разрешений изображений за пределами своей зоны обучения. Для преодоления этого ограничения мы концептуализируем изображения как последовательности токенов с динамическими размерами, вместо традиционных методов, которые воспринимают изображения как сетки фиксированного разрешения. Эта перспектива обеспечивает гибкую стратегию обучения, которая плавно адаптируется к различным соотношениям сторон как во время обучения, так и во время вывода, тем самым способствуя обобщению разрешения и устранению предвзятостей, внесенных обрезкой изображения. На основе этого мы представляем Гибкий Визионный Трансформер (FiT), архитектуру трансформера, специально разработанную для генерации изображений с неограниченными разрешениями и соотношениями сторон. Мы далее улучшаем FiT до FiTv2 с несколькими инновационными дизайнами, включая нормализацию векторов Запрос-Ключ, модуль AdaLN-LoRA, планировщик исправленного потока и выборку Логит-Нормал. Усиленная тщательно настроенной структурой сети, FiTv2 демонстрирует ускорение сходимости в 2 раза по сравнению с FiT. При использовании передовых техник экстраполяции без обучения FiTv2 проявляет замечательную адаптивность как в экстраполяции разрешения, так и в генерации разнообразных разрешений. Кроме того, наше исследование масштабируемости модели FiTv2 показывает, что более крупные модели обладают лучшей вычислительной эффективностью. Кроме того, мы представляем эффективную стратегию послеобучения для адаптации предварительно обученной модели для генерации высокого разрешения. Обширные эксперименты демонстрируют исключительную производительность FiTv2 на широком диапазоне разрешений. Мы опубликовали все коды и модели на https://github.com/whlzy/FiT для поощрения исследования моделей диффузионных трансформеров для генерации изображений произвольного разрешения.

English

Nature is infinitely resolution-free. In the context of this reality, existing diffusion models, such as Diffusion Transformers, often face challenges when processing image resolutions outside of their trained domain. To address this limitation, we conceptualize images as sequences of tokens with dynamic sizes, rather than traditional methods that perceive images as fixed-resolution grids. This perspective enables a flexible training strategy that seamlessly accommodates various aspect ratios during both training and inference, thus promoting resolution generalization and eliminating biases introduced by image cropping. On this basis, we present the Flexible Vision Transformer (FiT), a transformer architecture specifically designed for generating images with unrestricted resolutions and aspect ratios. We further upgrade the FiT to FiTv2 with several innovative designs, includingthe Query-Key vector normalization, the AdaLN-LoRA module, a rectified flow scheduler, and a Logit-Normal sampler. Enhanced by a meticulously adjusted network structure, FiTv2 exhibits 2times convergence speed of FiT. When incorporating advanced training-free extrapolation techniques, FiTv2 demonstrates remarkable adaptability in both resolution extrapolation and diverse resolution generation. Additionally, our exploration of the scalability of the FiTv2 model reveals that larger models exhibit better computational efficiency. Furthermore, we introduce an efficient post-training strategy to adapt a pre-trained model for the high-resolution generation. Comprehensive experiments demonstrate the exceptional performance of FiTv2 across a broad range of resolutions. We have released all the codes and models at https://github.com/whlzy/FiT to promote the exploration of diffusion transformer models for arbitrary-resolution image generation.