FiTv2: Skalierbarer und verbesserter flexibler Vision-Transformer für Diffusionsmodelle

FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model

October 17, 2024

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

cs.AI

Zusammenfassung

Die Natur ist unendlich frei von Auflösung. In diesem Zusammenhang stehen bestehende Diffusionsmodelle, wie Diffusions-Transformer, oft vor Herausforderungen bei der Verarbeitung von Bildauflösungen außerhalb ihres trainierten Bereichs. Um diese Einschränkung zu überwinden, konzeptualisieren wir Bilder als Sequenzen von Tokens mit dynamischen Größen, anstatt Bilder wie herkömmliche Methoden als fest aufgelöste Raster wahrzunehmen. Diese Perspektive ermöglicht eine flexible Schulungsstrategie, die während des Trainings und der Inferenz nahtlos verschiedene Seitenverhältnisse berücksichtigt, um so die Auflösungsverallgemeinerung zu fördern und Verzerrungen durch Bildbeschnitt zu eliminieren. Auf dieser Grundlage präsentieren wir den Flexiblen Visionstransformer (FiT), eine Transformer-Architektur, die speziell für die Erzeugung von Bildern mit uneingeschränkten Auflösungen und Seitenverhältnissen entwickelt wurde. Wir verbessern den FiT weiter zu FiTv2 mit mehreren innovativen Designs, einschließlich der Normalisierung des Query-Key-Vektors, des AdaLN-LoRA-Moduls, eines rektifizierten Flussplaners und eines Logit-Normal-Samplers. Durch eine sorgfältig angepasste Netzwerkstruktur verstärkt, zeigt FiTv2 eine 2-fache Konvergenzgeschwindigkeit im Vergleich zu FiT. Bei der Integration fortschrittlicher, trainingfreier Extrapolationstechniken zeigt FiTv2 eine bemerkenswerte Anpassungsfähigkeit sowohl bei der Auflösungsextrapolation als auch bei der Erzeugung verschiedener Auflösungen. Darüber hinaus zeigt unsere Untersuchung der Skalierbarkeit des FiTv2-Modells, dass größere Modelle eine bessere Rechenleistung aufweisen. Des Weiteren führen wir eine effiziente Post-Training-Strategie ein, um ein vortrainiertes Modell für die hochauflösende Erzeugung anzupassen. Umfangreiche Experimente zeigen die außergewöhnliche Leistung von FiTv2 über ein breites Spektrum von Auflösungen hinweg. Wir haben alle Codes und Modelle unter https://github.com/whlzy/FiT veröffentlicht, um die Erforschung von Diffusions-Transformer-Modellen für die Erzeugung von Bildern mit beliebiger Auflösung zu fördern.

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.