FiTv2: Vision Transformer Flessibile e Migliorato per la Diffusione Scalabile

FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model

October 17, 2024

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

cs.AI

Abstract

La natura è infinitamente priva di risoluzione. Nel contesto di questa realtà, i modelli di diffusione esistenti, come i Diffusion Transformers, spesso incontrano sfide nel processare risoluzioni delle immagini al di fuori del loro dominio addestrato. Per affrontare questa limitazione, concepiamo le immagini come sequenze di token con dimensioni dinamiche, piuttosto che i metodi tradizionali che percepiscono le immagini come griglie a risoluzione fissa. Questa prospettiva consente una strategia di addestramento flessibile che si adatta senza soluzione di continuità a vari rapporti d'aspetto sia durante l'addestramento che nell'inferenza, promuovendo così la generalizzazione della risoluzione ed eliminando i bias introdotti dal ritaglio delle immagini. Sulla base di questo, presentiamo il Flexible Vision Transformer (FiT), un'architettura transformer progettata specificamente per generare immagini con risoluzioni e rapporti d'aspetto illimitati. Potenziando ulteriormente il FiT al FiTv2 con diversi design innovativi, tra cui la normalizzazione del vettore Query-Key, il modulo AdaLN-LoRA, un programmatore di flusso rettificato e un campionatore Logit-Normal. Potenziato da una struttura di rete attentamente regolata, il FiTv2 mostra una velocità di convergenza 2 volte superiore al FiT. Integrando tecniche avanzate di estrapolazione senza addestramento, il FiTv2 dimostra una notevole adattabilità sia nell'estrapolazione della risoluzione che nella generazione di risoluzioni diverse. Inoltre, la nostra esplorazione della scalabilità del modello FiTv2 rivela che modelli più grandi mostrano una migliore efficienza computazionale. Inoltre, introduciamo una strategia efficiente di post-addestramento per adattare un modello pre-addestrato per la generazione ad alta risoluzione. Esperimenti esaustivi dimostrano le eccezionali prestazioni del FiTv2 su un'ampia gamma di risoluzioni. Abbiamo reso disponibili tutti i codici e i modelli su https://github.com/whlzy/FiT per promuovere l'esplorazione dei modelli di diffusione transformer per la generazione di immagini a risoluzione arbitraria.

English

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 2times 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.