FiTv2: Transformer de Visión Flexible Escalable y Mejorado para Modelos de Difusión

FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model

October 17, 2024

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

cs.AI

Resumen

La naturaleza es infinitamente libre de resolución. En el contexto de esta realidad, los modelos de difusión existentes, como los Transformadores de Difusión, a menudo enfrentan desafíos al procesar resoluciones de imágenes fuera de su dominio entrenado. Para abordar esta limitación, conceptualizamos las imágenes como secuencias de tokens con tamaños dinámicos, en lugar de los métodos tradicionales que perciben las imágenes como cuadrículas de resolución fija. Esta perspectiva permite una estrategia de entrenamiento flexible que acomoda sin problemas varias relaciones de aspecto durante tanto el entrenamiento como la inferencia, promoviendo así la generalización de resolución y eliminando sesgos introducidos por el recorte de imágenes. Sobre esta base, presentamos el Transformador de Visión Flexible (FiT), una arquitectura de transformer diseñada específicamente para generar imágenes con resoluciones y relaciones de aspecto ilimitadas. Mejoramos aún más el FiT a FiTv2 con varios diseños innovadores, incluyendo la normalización del vector de Consulta-Clave, el módulo AdaLN-LoRA, un programador de flujo rectificado y un muestreador Logit-Normal. Potenciado por una estructura de red meticulosamente ajustada, FiTv2 exhibe una velocidad de convergencia 2 veces mayor que FiT. Al incorporar técnicas avanzadas de extrapolación sin entrenamiento, FiTv2 demuestra una notable adaptabilidad tanto en la extrapolación de resolución como en la generación de resoluciones diversas. Además, nuestra exploración de la escalabilidad del modelo FiTv2 revela que los modelos más grandes muestran una mejor eficiencia computacional. Además, presentamos una estrategia eficiente de post-entrenamiento para adaptar un modelo pre-entrenado para la generación de alta resolución. Experimentos exhaustivos demuestran el rendimiento excepcional de FiTv2 en una amplia gama de resoluciones. Hemos publicado todos los códigos y modelos en https://github.com/whlzy/FiT para promover la exploración de modelos de transformer de difusión para la generación de imágenes de resolución 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.