FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model

TL;DR

This work conceptualizes images as sequences of tokens with dynamic sizes, rather than traditional methods that perceive images as fixed-resolution grids, which 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.

Abstract

\textit{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 \textbf{Flexible Vision Transformer} (FiT), a transformer architecture specifically designed for generating images with \textit{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 $2\times$ 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 \url{https://github.com/whlzy/FiT} to promote the exploration of diffusion transformer models for arbitrary-resolution image generation.

Figures (10)

• Figure 1: Selected samples from FiTv2-3B/2 models at resolutions of $256\times256$, $512\times512$, $768\times768$, $256\times768$ and $768\times256$. All the images are sampeld with CFG=4.0. FiT is capable of generating images at unrestricted resolutions and aspect ratios. FiTv2 pushes the image generation ability of FiT to a new level, capable of generating better and higher-resolution images.
• Figure 2: The Height/Width distribution of the original ImageNetdeng2009imagenet dataset.
• Figure 3: Overview of (a) flexible training pipeline, and (b) flexible inference pipeline. We conceptualize images as dynamic sequences of tokens, allowing for flexible image generation across different resolutions and aspect ratios.
• Figure 4: Block comparison between (a) FiT and (b) FiTv2. New modules, QKNorm, AdaLN-LoRA and Global AdaLN, are marked by red color.
• Figure 5: Pipeline comparison between (a) DiT, (b) FiT, and (c) FiTv2. In FiTv2, we incorporate both fixed-resolution images and the flexible-resolution images into training process.