FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model
Authors: ZiDong Wang, Zeyu Lu, Di Huang, Cai Zhou, Xiaoyu Yue, Wanli Ouyang, and Lei Bai
Organizations: Shanghai AI Laboratory, Shanghai, 200000, China. · Chinese University of Hong Kong, Shatin, 999077, Hong Kong. · Shanghai Jiao Tong University, Shanghai, 200000, China. · University of Sydney, Camperdown NSW 2050, Australia. · Tsinghua University, Beijing, 100084, China.
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 2× 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.
Figures & tables
Fig. 1: Selected samples from FiTv2-3B/2 models at resolutions of 256×256 , 512×512 , 768×768 , 256×768 and 768×256 . All the images are sampled 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.
Fig. 2: 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.
Fig. 3: Block comparison between (a) FiT and (b) FiTv2. New modules, QKNorm, AdaLN-LoRA and Global AdaLN, are marked by red color.
Fig. 4: 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.
Fig. 5: Effect of scaling FiTv2 model. All the images are sampled without using CFG. We demonstrate the FID score over training iterations (a) and training GFLOPs (b) of our FiTv2 model with three sizes. Scaling our FiTv2 model yields better quantitative and qualitative performance.
Fig. 6: Selected samples from FiTv2-XL/2 models at resolutions of 256×256 on text-to-image generation tasks. All the images are sampled with CFG=4.0. With only 400K training steps, our model is capable of generating releastic images according to text descriptions.
Fig. 7: Comparision of FID and CLIP-L score across different CFG scales for two text-to-image models: FiTv2-XL/2 and SiT-XL/2. FiTv2-XL/2 significantly outperforms SiT-XL/2 in terms of FID score and CLIP-L score.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Fig. 8: Uncurated samples from FiTv2-3B/2 models at resolutions of 512×512 , 256×768 and 768×256 .
Fig. 9: Uncurated samples from FiTv2-3B/2 models at resolutions of 512×512 , 256×768 and 768×256 .
Fig. 10: Uncurated samples from FiTv2-3B/2 models at resolutions of 512×512 , 256×768 and 768×256 .
Fig. 11: Uncurated samples from FiTv2-3B/2 models at resolutions of 512×512 , 256×768 and 768×256 .
Fig. 12: Uncurated samples from FiTv2-3B/2 models at resolutions of 512×512 , 256×768 and 768×256 .
Fig. 13: Uncurated samples from FiTv2-3B/2 models at resolutions of 512×512 , 256×768 and 768×256 .
Fig. 14: Uncurated samples from FiTv2-XL/2 models at resolutions of 256×256 on text-to-image generation tasks. All the images are sampled with CFG=4.0. With only 400K training steps, our model is capable of generating releastic images according to text descriptions.
Generating images beyond the training resolution of Diffusion Transformers (DiTs) remains challenging, as the number of tokens grows rapidly with resolution, full attention becomes expensive, and outputs often suffer from structural distortions and degraded detail. We identify uniform full attention as the common source of both problems. Attending to all tokens at every step causes quadratic cost and exposes the model to more tokens and relative distances beyond the training range, pushing the model out of distribution. We observe that attention's effective receptive field narrows during denoising, from global in early steps that form structure to local in later steps that refine details. We propose AReA, a training-free attention mechanism whose receptive field evolves from coarse to fine, following the denoising progression. In early steps, AReA applies cross-scale coarse attention over downsampled keys and values, capturing global context, while in later steps it applies local attention within progressively shrinking windows. Without additional parameters or training, experiments show that AReA consistently outperforms state-of-the-art training-free baselines across target resolutions and different models, while also being faster and more efficient.
Training-free text-to-high-resolution image generation has recently attracted growing research attention. However, existing studies on this task primarily focus on adapting off-the-shelf U-Net-based diffusion models to high resolutions, with limited progress on adapting off-the-shelf Diffusion Transformer (DiT) models despite their strong text-to-image generation capabilities at limited resolutions. In this work, we find two key challenges particularly hindering the application of off-the-shelf DiT models for high-resolution image synthesis in a training-free manner, namely, spatial disorder and long generation time. To address these challenges, we propose a novel method tailored to adapt off-the-shelf DiT models for high-resolution image synthesis. Extensive experiments show the efficacy of our method. Our code is available at: https://github.com/zylwithxy/HRDiT.
Yu Xue, Haoxuan Qu, Zhuoling Li +5
Lancaster University, Lancaster, UK · South China University of Technology, Guangzhou, China · NVIDIA AI Tech Centre, Singapore
Diffusion transformers (DiTs) have emerged as a dominant architecture for text-to-image generation, yet their performance drops when generating at resolutions beyond their training range. Existing training-free approaches mitigate this by modifying inference-time attention behavior, often through Rotary Position Embeddings (RoPE) extrapolation combined with attention scaling. However, these strategies apply a uniform and content-agnostic scaling across RoPE components with distinct frequency characteristics, inducing a trade-off between preserving global structure and recovering fine detail. We introduce SEGA, a training-free method that dynamically scales attention across RoPE components according to the latent's spatial-frequency structure at each denoising step. This adaptive scaling improves both structural coherence and fine-detail fidelity. Experiments show that SEGA consistently improves high-resolution synthesis across multiple target resolutions, outperforming state-of-the-art training-free baselines.