PaperlayerView exact revision on arXiv

Source-linked AI summary

FiT: Flexible Vision Transformer for Diffusion Model

Zeyu Lu, Zidong Wang, Di Huang, Chengyue Wu, Xihui Liu, Wanli Ouyang, Lei Bai

arXiv:2402.12376v4cs.CV

TL;DR

Diffusion transformers such as DiT struggle with resolutions outside their trained range because they cannot use dynamic-resolution images during training. FiT models images as variable-length token sequences, preserves aspect ratios during flexible training, modifies the transformer architecture, and applies training-free resolution extrapolation. After 1.8 million steps, FiT-XL/2 outperforms state-of-the-art CNN and transformer models at five listed resolutions, but slightly trails DiT-XL/2 at 256 × 256.

  • Problem

    Diffusion transformers such as DiT struggle with resolutions outside their trained range because they cannot use dynamic-resolution images during training.

    Source: §1 · p. 2

  • Method

    FiT models images as variable-length token sequences, preserves aspect ratios during flexible training, modifies the transformer architecture, and applies training-free resolution extrapolation.

    Source: §1 · p. 2, §3.4 · p. 5

  • Results

    After 1.8 million steps, FiT-XL/2 outperforms state-of-the-art CNN and transformer models at five listed resolutions, but slightly trails DiT-XL/2 at 256 × 256.

    Source: §1 · p. 2

  • Takeaways & Limitations

    FiT is presented as a flexible diffusion transformer for image generation at arbitrary resolutions and aspect ratios, with VisionNTK further enhancing performance.

    Source: §5 · p. 8

    With only 1800K training steps, FiT-XL/2 is slightly inferior to DiT-XL/2 at 256x256; its generative capabilities with higher-resolution training have not yet been thoroughly explored.

    Source: C. Limitations and Future Work · p. 12

Abstract

from arXiv · show

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 overcome this limitation, we present the Flexible Vision Transformer (FiT), a transformer architecture specifically designed for generating images with unrestricted resolutions and aspect ratios. Unlike traditional methods that perceive images as static-resolution grids, FiT conceptualizes images as sequences of dynamically-sized tokens. This perspective enables a flexible training strategy that effortlessly adapts to diverse aspect ratios during both training and inference phases, thus promoting resolution generalization and eliminating biases induced by image cropping. Enhanced by a meticulously adjusted network structure and the integration of training-free extrapolation techniques, FiT exhibits remarkable flexibility in resolution extrapolation generation. Comprehensive experiments demonstrate the exceptional performance of FiT across a broad range of resolutions, showcasing its effectiveness both within and beyond its training resolution distribution. Repository available at https://github.com/whlzy/FiT.

1. Introduction

FiT addresses resolution limits in diffusion transformers by modeling images as variable-length token sequences and adapting its training, architecture, and inference pipeline. Its FiT-XL/2 model reports strong results across several resolutions, with performance further improved by training-free extrapolation.

1. Introduction

  • DiT struggles with arbitrary resolutions because it cannot use dynamic-resolution images during training, limiting adaptation to different token lengths and resolutions.This is the motivation for FiT's flexible image modeling approach.Source: §1 · p. 2
  • FiT treats images as variable-length token sequences, padding them to a maximum specified length to support generation at unrestricted resolutions and aspect ratios.Its design spans a flexible training pipeline, network architecture, and inference process.Source: §1 · p. 2
  • FiT's training pipeline preserves each image's aspect ratio and resizes high-resolution images to a maximum token limit without cropping or disproportionate scaling.The authors describe it as the first transformer-based generation model to maintain diverse image resolutions throughout training.Source: §1 · p. 2
  • FiT evolves from DiT with 2D RoPE, SwiGLU in place of the MLP, and Masked MHSA to manage padding tokens.The passage notes 2D RoPE's inspiration from length extrapolation in large language models.Source: §1 · p. 2
  • Directly applying language-model token-length extrapolation techniques to FiT performs suboptimally, so the authors tailor them for 2D RoPE.The tailored techniques enhance performance across resolutions and aspect ratios.Source: §1 · p. 2
  • After 1.8 million training steps, FiT-XL/2 outperforms state-of-the-art CNN and transformer models at five listed resolutions, while slightly trailing DiT-XL/2 at 256 × 256.The listed resolutions are 160×320, 128×384, 320×320, 224×448, and 160×480; training-free resolution extrapolation further improves FiT-XL/2.Source: §1 · p. 2

2. Related Work

Related work spans diffusion generation, transformer backbones for vision and latent diffusion, and sequence-length extrapolation in language models. FiT follows the latent-diffusion transformer line and studies RoPE extrapolation for vision generation.

2. Related Work

  • LDMs (Rombach et al., 2022) reverse noise in latent space using VQ-VAE (Esser et al., 2021), while DDIM (Song et al., 2020a) offers accelerated sampling.The passage also describes progress in DDPMs and score-based image-generation models.Source: §2 · pp. 2–3
  • DiT (Peebles & Xie, 2023) uses a vision transformer as an LDM backbone and serves as a strong baseline; FiT follows this LDM paradigm.Other cited extensions include MDT's masked latent modeling and U-ViT's integration of U-Net architectures into a ViT backbone.Source: §2 · p. 3
  • RoPE allows flexible sequence lengths but loses performance beyond training length; PI (Chen et al., 2023) and NTK-aware (LocalLLaMA) are among proposed extension methods.The paper studies RoPE implementation and on-the-fly resolution extrapolation for vision generation.Source: §2 · p. 3

3. Flexible Vision Transformer for Diffusion

FiT treats images as variable-length token sequences, combining flexible training, padding-aware attention, and 2D positional embeddings with training-free interpolation for resolution extrapolation.

3.1. Preliminary

  • 1-D RoPE unifies absolute and relative positional information; NTK-aware interpolation modifies its rotary base, while YaRN adjusts rotary frequencies using a dimension-dependent ramp.The cited methods address extrapolation to context lengths beyond training length.Source: §3.1 · p. 3

3.2. Flexible Training and Inference Pipeline

  • Training avoids cropping and low-resolution upscaling, resizing only high-resolution images to a maximum area of 256².Variable-length latent-token sequences are padded to Lmax=256 for batching, and loss is calculated only for denoised output tokens.Source: §3.2 · p. 4
  • At inference, FiT samples noisy tokens, denoises them for K iterations, then reshapes and unpatchifies them according to a predefined position map.Source: §3.2 · p. 4

3.3. Flexible Vision Transformer Architecture

  • Masked MHSA lets noised tokens attend while blocking interactions with padding tokens, which standard DiT attention cannot distinguish.Noised tokens receive mask value 0 and padding tokens receive negative infinity.Source: §3.3 · p. 4
  • 2D-RoPE computes positional rotations separately for height and width coordinates, then concatenates them to support resolution generalization.Its attention score has no height-width cross-term, allowing rotary frequencies to be decoupled.Source: §3.3 · pp. 4–5
  • FiT replaces the FFN's traditional MLP with SwiGLU.SwiGLU multiplies SiLU(xW) and xV elementwise.Source: §3.3 · p. 5

3.4. Training Free Resolution Extrapolation

  • VisionNTK and VisionYaRN are training-free positional interpolation methods for out-of-distribution positions, especially effective when generating images with arbitrary aspect ratios.At aspect ratio one, they are equivalent to vanilla NTK and YaRN implementations.Source: §3.4 · p. 5

4. Experiments

FiT is evaluated on ImageNet across in- and out-of-distribution resolutions and aspect ratios, with experiments testing its training pipeline, architecture, and extrapolation methods. FiT-XL/2 performs strongly across these settings, while ablations examine how flexible training and architectural choices contribute.

4.1. FiT Implementation

  • The experiments use class-conditional latent FiT models trained with a maximum of 256 tokens and assess multiple aspect ratios and token lengths inside and beyond that limit.FID-50K is the major metric; sFID, IS, Precision, and Recall are also reported.Source: §4.1 · pp. 5–6

4.2. FiT Architecture Design

  • For DiT-B/2, Config A with flexible training alone reduces FID by 40.81 at 160 × 320 and 56.55 at 224 × 448, compared with fixed-resolution training.At 256 × 256, it improves FID by 1.49 versus DiT-B/2 with fixed-resolution training.Source: §4.2 · p. 7
  • Replacing MLP with SwiGLU reduces FID by 1.59, 1.85, and 0.21 at 256×256, 160×320, and 224×448, respectively.These comparisons are between Config B and Config A.Source: §4.2 · p. 7, Table 3 · p. 7
  • Using only 2D RoPE instead of absolute positional encoding reduces FID by 6.05, 5.45, and 6.39 across the three tested resolutions.Adding absolute positional encoding alongside 2D RoPE worsens FID by 1.82, 1.65, and 0.44 compared with using 2D RoPE alone.Source: §4.2 · p. 7, Table 3 · p. 7
  • FiT-B/2 reduces DiT-B/2 FID by 8.47 at 256 × 256 and by 47.36 and 64.43 at 160×320 and 224×448.The full-model comparison is reported at 400K training steps.Source: §4.2 · p. 7, Table 3 · p. 7

4.3. FiT Resolution Extrapolation Design

  • VisionYaRN reduces FID versus YaRN by 40.27 at 224 × 448 and 41.22 at 160 × 480, while VisionNTK outperforms direct extrapolation at all three tested resolutions.Vanilla YaRN improves 320 × 320 but worsens the other two resolutions; PI decreases FiT-B/2 performance at all three.Source: Table 4 · p. 7, §4.3 · p. 8

4.4. FiT In-Distribution Resolution Results

  • FiT-XL/2 reduces prior-best FID-50K from 6.93 to 5.74 at 160 × 320 and from 29.67 to 16.81 at 128 × 384.At 256 × 256, its FID is slightly higher than models trained for longer.Source: Table 1 · p. 6, §4.4 · p. 8

4.5. FiT Out-Of-Distribution Resolution Results

  • FiT-XL/2 achieves the best FID-50K and IS at all three out-of-distribution resolutions, exceeding prior-best LDM FID-50K by 0.82, 0.65, and 3.52.The evaluated resolutions are 320×320, 224×448, and 160×480; sFID performance is described as competitive.Source: Table 2 · p. 6, §4.5 · p. 8

5. Conclusion

FiT is a refined transformer architecture with a flexible training pipeline for generating images at arbitrary resolutions and aspect ratios. It surpasses previous transformer-based and CNN-based models across various resolutions, with VisionNTK further enhancing performance.

  • FiT is a refined transformer architecture with a flexible training pipeline designed for generating images at arbitrary resolutions and aspect ratios.Source: §5 · p. 8
  • FiT surpasses previous transformer-based and CNN-based models across various resolutions.Source: §5 · p. 8
  • VisionNTK significantly enhances FiT's performance further.Source: §5 · p. 8

A. Experimentin Setups

The appendix reports where model configurations and performance are listed, and specifies the shared image encoding and decoding setup and evaluation suite.

B. Detailed Attention Score with 2D RoPE and decoupled 2D-RoPE.

The appendix defines 2D RoPE as a vector-valued complex function and describes its use in attention scores. It also explains that the height and width rotary frequencies can be decoupled because the attention score has no cross-term between them.

C. Limitations and Future Work

The authors identify compute-limited training and an incomplete exploration of both higher-resolution training and extrapolation methods as boundaries of their study.

D. More Model Samples

This section presents uncurated FiT-XL/2 samples at square and non-square resolutions, alongside failure samples from DiT-XL/2.

D. More Model Samples

Related papers

Loading 2402.12376v4…