FiTv2: Transformer de Visão Flexível Aprimorado e Escalável para Modelos de Difusão

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

Resumo

A natureza é infinitamente livre de resolução. No contexto desta realidade, os modelos de difusão existentes, como os Transformadores de Difusão, frequentemente enfrentam desafios ao processar resoluções de imagem fora de seu domínio treinado. Para lidar com essa limitação, conceitualizamos imagens como sequências de tokens com tamanhos dinâmicos, em vez de métodos tradicionais que percebem imagens como grades de resolução fixa. Essa perspectiva possibilita uma estratégia de treinamento flexível que acomoda facilmente várias razões de aspecto durante o treinamento e a inferência, promovendo assim a generalização de resolução e eliminando vieses introduzidos pelo recorte de imagem. Com base nisso, apresentamos o Transformador de Visão Flexível (FiT), uma arquitetura de transformer projetada especificamente para gerar imagens com resoluções e razões de aspecto ilimitadas. Aprimoramos ainda mais o FiT para o FiTv2 com vários designs inovadores, incluindo a normalização do vetor de Consulta-Chave, o módulo AdaLN-LoRA, um agendador de fluxo retificado e um amostrador Logit-Normal. Aprimorado por uma estrutura de rede meticulosamente ajustada, o FiTv2 exibe uma velocidade de convergência 2 vezes maior que o FiT. Ao incorporar técnicas avançadas de extrapolação sem treinamento, o FiTv2 demonstra uma notável adaptabilidade tanto na extrapolação de resolução quanto na geração de resoluções diversas. Além disso, nossa exploração da escalabilidade do modelo FiTv2 revela que modelos maiores apresentam uma melhor eficiência computacional. Além disso, introduzimos uma estratégia eficiente de pós-treinamento para adaptar um modelo pré-treinado para a geração de alta resolução. Experimentos abrangentes demonstram o desempenho excepcional do FiTv2 em uma ampla gama de resoluções. Disponibilizamos todos os códigos e modelos em https://github.com/whlzy/FiT para promover a exploração de modelos de transformadores de difusão para geração de imagens de resolução arbitrária.

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.