Research map / Image generation / Autoregressive and latent generation
FiT: Flexible Vision Transformer for Diffusion Model
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes FiT, a diffusion transformer that trains on native-aspect-ratio images packed as variable-length token sequences with masked attention and 2D RoPE, plus training-free VisionNTK/VisionYaRN extrapolation for arbitrary resolutions.
Problem
Diffusion transformers such as DiT are trained on fixed-resolution, center-cropped square images and generalize poorly to other resolutions and aspect ratios.
Approach
Treat images as sequences of dynamically sized tokens: preserve aspect ratio without cropping, pad to a max token length with masked attention, replace absolute PE with 2D RoPE, and adapt NTK/YaRN interpolation separately per height and width for training-free extrapolation.
Key results
- FiT-XL/2 with VisionNTK reaches FID-50K 5.74 at 160x320 (vs. U-ViT-H/2-G 6.93) and 16.81 at 128x384 (vs. LDM-4-G 29.67) at in-distribution resolutions.
- At out-of-distribution resolutions, FiT-XL/2 beats the previous best LDM-4 by 0.82, 0.65 and 3.52 FID at 320x320, 224x448 and 160x480 respectively.
- On FiT-B/2 at 400K steps, VisionYaRN lowers FID at 224x448 from 82.19 (vanilla YaRN) to 41.92 and VisionNTK reaches 43.84 vs. DiT-B/2 at 94.94-109.1 range.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| train with flexible resolutions instead of fixed-resolution training (ablation) | DiT-B/2 · ImageNet 160×320 | FID: 91.32 -> 50.51 | Helped |
| train with flexible resolutions instead of fixed-resolution training (ablation) | DiT-B/2 · ImageNet 224×448 | FID: 109.1 -> 52.55 | Helped |
| replace MLP with SwiGLU (ablation) | flexibly trained DiT-B/2 · ImageNet 256×256, 160×320, and 224×448 | FID: 43.34 -> 41.75; 50.51 -> 48.66; 52.55 -> 52.34 | Helped |
| replace absolute PE with 2D RoPE (ablation) | flexibly trained DiT-B/2 · ImageNet 256×256, 160×320, and 224×448 | FID: 43.34 -> 37.29; 50.51 -> 45.06; 52.55 -> 46.16 | Helped |
| retain absolute PE in addition to 2D RoPE (ablation) | flexibly trained DiT-B/2 · ImageNet 256×256, 160×320, and 224×448 | FID: 37.29 -> 39.11; 45.06 -> 46.71; 46.16 -> 46.60 | Hurt |
| use flexible training, 2D RoPE, and SwiGLU (ablation) | DiT-B/2 · ImageNet 256×256, 160×320, and 224×448 | FID: 44.83 -> 36.36; 91.32 -> 43.96; 109.1 -> 44.67 | Helped |
| apply PI (ablation) | FiT-B/2 · ImageNet 320×320, 224×448, and 160×480 | FID: 61.35 -> 65.76; 44.67 -> 175.42; 56.81 -> 224.83 | Hurt |
| apply vanilla YaRN (ablation) | FiT-B/2 · ImageNet 320×320, 224×448, and 160×480 | FID: 61.35 -> 44.76; 44.67 -> 82.19; 56.81 -> 104.06 | Mixed |
| replace vanilla YaRN with VisionYaRN (ablation) | FiT-B/2 · ImageNet 224×448 and 160×480 | FID: 82.19 -> 41.92; 104.06 -> 62.84 | Helped |
| use VisionNTK instead of direct extrapolation (ablation) | FiT-B/2 · ImageNet 320×320, 224×448, and 160×480 | FID: 61.35 -> 57.31; 44.67 -> 43.84; 56.81 -> 56.76 | Helped |
| use FiT-XL/2 with VisionNTK rather than the best listed baseline | class-conditional ImageNet generation · ImageNet 160×320 | FID: 6.93 -> 5.74 | Helped |
| use FiT-XL/2 with VisionNTK rather than the best listed baseline | class-conditional ImageNet generation · ImageNet 128×384 | FID: 29.67 -> 16.81 | Helped |
Surprising findings
- Direct PI interpolation on FiT-B/2 degrades performance at all tested out-of-distribution resolutions, and vanilla YaRN helps only at 320x320.
- FiT-XL/2 is slightly worse than DiT-XL/2 at the 256x256 training resolution (FID 4.27 vs 2.27 with CFG).
Limitations
- FiT-XL/2 trained only 1.8M steps due to limited compute and slightly underperforms DiT-XL/2 at 256x256.
- Higher token-length limits (larger max resolution) were not explored.
- Only training-free extrapolation methods were studied; training-based extrapolation was not explored.
Experimental setup
Datasets: ImageNet
Benchmarks: ImageNet
Metrics: FID-50K sFID IS Precision Recall
Baselines: DiT-XL/2 U-ViT-H/2 ADM LDM-4 MDT BigGAN-deep StyleGAN-XL MaskGIT CDM
Builds on: DiT RoPE NTK-aware interpolation YaRN SwiGLU
Closest papers by meaning
- LaVin-DiT: Large Vision Diffusion Transformer (CVPR 2025 · similarity 0.88)
- Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models (CVPR 2025 · similarity 0.88)
- Simpler Diffusion: 1.5 FID on ImageNet512 with Pixel-space Diffusion (CVPR 2025 · similarity 0.87)
- FlexiDiT: Your Diffusion Transformer Can Easily Generate High-Quality Samples with Less Compute (CVPR 2025 · similarity 0.87)
- ScaleCrafter: Tuning-free Higher-Resolution Visual Generation with Diffusion Models (ICLR 2024 · similarity 0.87)
- Alleviating Distortion in Image Generation via Multi-Resolution Diffusion Models and Time-Dependent Layer Normalization (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.