FiTv2 : Vision Transformer Flexible amélioré et évolutif pour la diffusion de modèles

FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model

October 17, 2024

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

cs.AI

Résumé

La nature est infiniment exempte de résolution. Dans le contexte de cette réalité, les modèles de diffusion existants, tels que les Transformateurs de Diffusion, rencontrent souvent des défis lors du traitement de résolutions d'images en dehors de leur domaine d'entraînement. Pour remédier à cette limitation, nous conceptualisons les images comme des séquences de jetons de tailles dynamiques, plutôt que les méthodes traditionnelles qui perçoivent les images comme des grilles de résolution fixe. Cette perspective permet une stratégie d'entraînement flexible qui s'adapte facilement à divers rapports hauteur/largeur à la fois pendant l'entraînement et l'inférence, favorisant ainsi la généralisation de la résolution et éliminant les biais introduits par le recadrage d'images. Sur cette base, nous présentons le Transformateur de Vision Flexible (FiT), une architecture de transformer spécifiquement conçue pour générer des images avec des résolutions et des rapports hauteur/largeur illimités. Nous améliorons ensuite le FiT en FiTv2 avec plusieurs conceptions innovantes, notamment la normalisation des vecteurs de requête-clé, le module AdaLN-LoRA, un planificateur de flux rectifié et un échantillonneur Logit-Normal. Renforcé par une structure de réseau méticuleusement ajustée, le FiTv2 présente une vitesse de convergence 2 fois supérieure à celle du FiT. En incorporant des techniques d'extrapolation avancées sans entraînement, le FiTv2 démontre une adaptabilité remarquable à la fois dans l'extrapolation de résolution et la génération de résolutions diverses. De plus, notre exploration de la scalabilité du modèle FiTv2 révèle que les modèles plus grands présentent une meilleure efficacité computationnelle. En outre, nous introduisons une stratégie efficace de post-entraînement pour adapter un modèle pré-entraîné à la génération haute résolution. Des expériences complètes démontrent les performances exceptionnelles du FiTv2 sur une large gamme de résolutions. Nous avons mis à disposition tous les codes et modèles sur https://github.com/whlzy/FiT pour promouvoir l'exploration des modèles de transformateur de diffusion pour la génération d'images à résolution arbitraire.

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.