FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model

arXiv:2410.13925 · cs.LG · Submitted 2024-10-17 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model".

Jane: Here's the gist of the paper. The gist The Flexible Vision Transformer (FiTv2) is a transformer architecture designed to generate images with unrestricted resolutions and aspect ratios by conceptualizing images as sequences of tokens with dynamic sizes, which enables flexible training and inference across various resolutions How it works The core motivation for FiT is the fundamental shift in image data conceptualization where FiT "conceptualizes images as sequences of variable-length tokens, departing from the traditional perspective of static grids with fixed dimensions".

Tom: This paradigm shift allows for dynamic adjustment of sequence length, which facilitates image generation at arbitrary resolutions unconstrained by predefined grids.

Jane: FiT achieves resolution-independent image synthesis by efficiently managing and padding these variable-length token sequences to a specified maximum.

Tom: FiT introduces three significant advancements to realize this paradigm shift: 1.

Paper summary: Jane: Moving on, let's look at the conclusion of "FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model". The paper really hammers home that this framework moves beyond just making things work on specific sizes; it addresses the fundamental limitation where existing diffusion models fail outside their trained domain.

Lu: They emphasize that by treating images as sequences of tokens with dynamic sizes, you get a flexible training strategy that lets you accommodate various aspect ratios during both training and inference without having to deal with the biases introduced by image cropping. This is a big shift from the traditional fixed-resolution grid idea.

Tom: So, the thesis is that this token-based approach enables resolution generalization across different sizes, which means if you train it on one set of images, it performs better when generating images at resolutions you didn't explicitly see during training.

Meng: But what about the practical implications for someone actually trying to deploy this? If we use a model that can handle unrestricted resolutions, how does that translate to our operational costs or the complexity of managing different generation requests?

Jane: The summary points out that they propose several innovative designs in FiTv2, including Query-Key vector normalization, the AdaLN-LoRA module for efficiency, and a rectified flow scheduler. These aren't just cosmetic changes; they are meant to make the system more stable and scalable.

Lu: Specifically, the paper details how these modifications lead to FiTv2 being a more efficient and scalable version of FiT by adjusting hidden sizes in the Swish-Gated Linear Unit and using a global AdaLN module to capture overlapping condition information.

Lalam: I see that the authors are also looking at training-free resolution extrapolation methods, which is pretty ambitious. They are benchmarking its performance against other state-of-the-art methods to show that this approach actually generalizes well beyond what was seen in the original training steps>Page three> <ref:2410.13925#pg2>

Tom: That's a lot of moving parts. We see claims about surpassing models like DiT two-XL/two and SiT eight-XL/two with the same parameters and only twenty-eight point six percent of the training costs on page two which is a significant point when you look at efficiency metrics <ref:2410.13925#pg2,DiT 2 -XL/2 and SiT 8 -XL/2 with the same>.

Jane: The authors also conduct an efficient post-training experiment to investigate how well it transfers from low resolution to high resolution, which shows they've addressed the gap between their standard benchmarks and real high-resolution generation.

Meng: So, the summary is that this work provides a more efficient and scalable version of FiT that handles resolutions flexibly, but I wonder if those efficiency gains hold up when we look at real-world traffic where things might break at scale.

Lu: The key idea they are pushing is using these dynamic token sequences to bypass the grid limitations entirely, which is what allows for the unrestricted generation capability they claim>Page one> <ref:2410.13925#pg1>

Tom: It’s a technical direction that aims to eliminate resolution biases introduced by cropping and offers a more flexible way to train these diffusion models. Now we've got a better idea of what FiTv2 is supposed to be achieving.

Conclusion: Jane: So, wrapping up this discussion on "FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model", the authors are presenting an architecture that fundamentally changes how we conceptualize image data by using dynamic tokens.

Lu: They successfully extended FiT into FiTv2 by incorporating specific innovations like Query-Key vector normalization and the AdaLN-LoRA module to make it more stable and parameter efficient. This results in a model capable of generating images at unrestricted resolutions and aspect ratios.

Tom: The implication here, based on what we've heard, is that we can expect diffusion models to be much better at handling varied input sizes during training and generation without performance degradation due to resolution mismatch.

Meng: For deployment, this means potentially less need for complex preprocessing pipelines where you have to resize or crop images before feeding them into the model, which would simplify the entire workflow significantly.

Jane: It’s about making the training process more flexible so that when we do use these models in production, they are inherently better at generalizing to unexpected image sizes. The paper shows that FiTv2 can push image generation ability past previous state-of-the-art models like SiT eight-XL/two on text-to-image tasks>Page two> <ref:2410.13925#pg2,on text-to-image tasks>

Lu: Ultimately, the paper suggests that this sequence token approach is a viable way to achieve resolution independence for vision transformers in the context of diffusion models.

Tom: So,FiTv2 is presented as an improved architecture that offers superior performance and scalability compared to FiT while maintaining its core flexible design principles. We've seen claims about faster convergence and better efficiency metrics from the paper>Page two> <ref:2410.13925#pg2>

Jane: That’s the gist of it—it’s a refined transformer designed to generate images without being tied down by fixed grid constraints, allowing for much broader use cases in image synthesis.

ZiDong Wang, Zeyu Lu, Di Huang, Cai Zhou, Wanli Ouyang, Lei Bai

Shanghai AI Laboratory · Chinese University of Hong Kong · Shanghai Jiao Tong University

cs.LG

Submitted: 2024-10-17

Updated: 2026-10-20

Comments: arXiv admin note: text overlap with arXiv:2402.12376

Code: https://github.com/whlzy/FiT

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 77/100

The gist: The gist The Flexible Vision Transformer (FiTv2) is a transformer architecture designed to generate images with unrestricted resolutions and aspect ratios by conceptualizing images as sequences of

Key concepts

Flexible Vision Transformer (FiT): FiT redefines image processing by viewing images as sequences of tokens with dynamic sizes instead of static grids. This allows the model to handle various resolutions and aspect ratios during both training and inference without needing to be specifically trained for one fixed size, promoting resolution generalization.
Query-Key vector normalization (QK-Norm): This technique stabilizes the attention mechanism by applying LayerNorm to the Query (Q) and Key (K) vectors before calculating attention. This effectively eliminates excessively large values in the attention logits, which significantly improves training stability, especially when using mixed-precision training methods.
AdaLN-LoRA module: Inspired by W.A.L.T., this module reduces parameter count by utilizing a global AdaLN structure to capture overlapping condition information across blocks. It is combined with LoRA to further reduce parameters, allowing for better performance while maintaining computational efficiency.
Logit-Normal sampler: This sampling method accelerates model convergence during the training schedule. By defining the sampling variable using a specific mathematical transformation, it helps mitigate challenges in learning velocity in the middle of the diffusion process, leading to faster and more efficient training.

Terminology

Transformer — 5× in this paper · explained in 2 episodes
Diffusion Model — 2× in this paper · explained in 10 episodes
Adaptability — 1× in this paper · explained in 1 episode
Diffusion Models — 1× in this paper · explained in 9 episodes
Generalization — 1× in this paper · explained in 13 episodes
Rectified Flow — 1× in this paper · explained in 3 episodes
Transformer Architecture — 1× in this paper · explained in 3 episodes
Transformer Model — 1× in this paper · explained in 2 episodes
Transformer Models — 1× in this paper · explained in 2 episodes
Transformers — 1× in this paper · explained in 2 episodes

Summary

The gist The Flexible Vision Transformer (FiTv2) is a transformer architecture designed to generate images with unrestricted resolutions and aspect ratios by conceptualizing images as sequences of tokens with dynamic sizes, which enables flexible training and inference across various resolutions

How it works

The core motivation for FiT is the fundamental shift in image data conceptualization where FiT conceptualizes images as sequences of variable-length tokens, departing from the traditional perspective of static grids with fixed dimensions. This paradigm shift allows for dynamic adjustment of sequence length, which facilitates image generation at arbitrary resolutions unconstrained by predefined grids. FiT achieves resolution-independent image synthesis by efficiently managing and padding these variable-length token sequences to a specified maximum.

FiT introduces three significant advancements to realize this paradigm shift:

A flexible training pipeline that preserves original image aspect ratios by treating images as token sequences, accommodating varied resolutions within a predefined maximum token limit.
A network architecture that incorporates 2-D Rotary Positional Embedding (2-D RoPE) [4], SwishGated Linear Units (SwiGLU) [5], and Masked Multi-Head Self-Attention to effectively handle diverse image sizes.
An inference process that adapts length extrapolation techniques from LLMs, tailoring them for 2-D RoPE to enhance performance across a wide range of resolutions and aspect ratios.

Enhancements in FiTv2

FiTv2 is an improved version of FiT built upon several innovative designs to achieve better performance and scalability:

**: Query-Key vector normalization (QK-Norm) was incorporated into the attention layer for stability, as it effectively eliminate[s] excessively large values in attention logits, which stabilizes the training process, particularly during mixed-precision training. This is achieved by applying LayerNorm to the Query (Q) and Key (K) vectors before the attention calculation. The hidden size of SwiGLU in FiTv2 is set to 2/3× that in the original FiT to align parameters and FLOPs with the baseline MLP. The parameter ratio of Attention, SwiGLU, and AdaLN module is 2: 4: 3. The use of AdaLN-LoRA is inspired by W.A.L.T. [48] to reduce the parameters occupied by the AdaLN module. A global AdaLN module is utilized to capture overlapping condition information and reduce the redundancy of condition information of each block. The final parameter composition for a FiTv2 block is N = Nattn + Nswiglu + NAdaLN-LoRA = 4 · d2 + 8 · d2 + 1.75 · d2. The flexibility allows reducing the LoRA rank r while simultaneously increasing the number of attention layers N, leading to enhanced model performance. The Logit-Normal sampler is adopted, which facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. A novel mixed data preprocessing strategy was proposed to bridge the gap between FiT's performance on standard benchmarks and its ability to handle diverse resolutions. The Logit-Normal sampler is defined as u ∼ N (0, 1), t = log(u(1 − u)). The mixed data preprocessing strategy involves randomly selecting between resizing only or resizing and cropping with a probability of 1/2 for images whose sizes are both larger than the target resolution size. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy mitigates blurring and information loss issues prevalent in previous methods. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence. The mixed data preprocessing strategy aligns our generation distribution with that of the ADM ImageNet reference dataset used for our FID evaluation. The Logit-Normal sampler facilitates the challenge of learning velocity in the middle of the schedule, as highlighted by EDM [50], and significantly accelerates the model convergence.

Improvements for AI systems

Enhanced resolution generalization by conceptualizing images as sequences of tokens with dynamic sizes rather than fixed-resolution grids, enabling a flexible training strategy that seamlessly accommodates various aspect ratios during both training and inference. This allows the system to generate images at unrestricted resolutions and aspect ratios without relying on image cropping biases.
Improved training stability and convergence speed by incorporating Query-Key vector normalization (QK-Norm) into the attention layer, which is described as contributing to stabilizing the training process. Furthermore, switching from DDPM to rectified flow and adopting the LogitNormal sampler results in 2x convergence speed of FiT, making training significantly faster.
Superior high-resolution image generation by implementing a mixed data preprocessing strategy that randomly selects between resizing only or resizing and cropping for larger images, which mitigates the blurring and information loss issues prevalent in previous methods and enables the model to achieve competitive performance on standard ImageNet benchmarks.
Increased model efficiency during training by adopting the AdaLN-LoRA module instead of a full SwiGLU MLP, leading to a parameter ratio of 2 : 4 : 3 for Attention, SwiGLU, and AdaLN modules and resulting in a decrease in model parameters occupied by the AdaLN module.
Robust extrapolation capabilities across diverse resolutions by utilizing VisionNTK combined with an attention scale factor defined as s = max(1.0, q log Htest×Wtest / Htrain×Wtrain), which significantly surpasses all the other extrapolation methods when generating images out of distribution on high-resolution aspect ratios.
Superior text-to-image generation performance by leveraging the 2D RoPE and a modulation mechanism of AdaLN for text conditioning, which involves averaging-pooling the text tokens into a semantic text embedding cT to replace the original class embedding.

Abstract

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 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 https://github.com/whlzy/FiT to promote the exploration of diffusion transformer models for arbitrary-resolution image generation.