FiTv2: 확장 가능하고 향상된 유연한 비전 트랜스포머(Diffusion 모델)

FiTv2: Scalable and Improved Flexible Vision Transformer for Diffusion Model

October 17, 2024

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

cs.AI

초록

자연은 무한한 해상도를 가지고 있습니다. 이 현실 속에서 기존의 확산 모델인 확산 트랜스포머와 같은 모델은 종종 훈련된 도메인을 벗어나는 이미지 해상도를 처리할 때 어려움을 겪습니다. 이 한계를 극복하기 위해 이미지를 고정 해상도 그리드로 인식하는 전통적인 방법이 아닌 동적 크기의 토큰 시퀀스로 개념화하여 이미지를 다룹니다. 이 관점은 훈련 및 추론 중에 다양한 종횡비를 유연하게 수용할 수 있는 훈련 전략을 제공하여 해상도 일반화를 촉진하고 이미지 자르기로 인한 편향을 제거합니다. 이를 기반으로, 제한 없는 해상도와 종횡비로 이미지를 생성하기 위해 특별히 설계된 트랜스포머 아키텍처인 유연 비전 트랜스포머(FiT)를 제안합니다. 우리는 쿼리-키 벡터 정규화, AdaLN-LoRA 모듈, 교정된 플로우 스케줄러, 로짓-정규 샘플러를 포함한 여러 혁신적인 설계를 통해 FiT를 FiTv2로 업그레이드합니다. 정교하게 조정된 네트워크 구조에 의해 강화된 FiTv2는 FiT의 수렴 속도의 2배를 보여줍니다. 고급 훈련 없는 외삽 기술을 통합할 때, FiTv2는 해상도 외삽과 다양한 해상도 생성 모두에서 놀라운 적응성을 보여줍니다. 또한, FiTv2 모델의 확장성 탐색을 통해 더 큰 모델이 더 나은 계산 효율성을 보인다는 것을 밝혀냅니다. 더불어, 미리 훈련된 모델을 고해상도 생성에 적응시키기 위한 효율적인 후훈련 전략을 소개합니다. 포괄적인 실험을 통해 FiTv2의 우수한 성능을 다양한 해상도 범위에서 입증했습니다. 임의 해상도 이미지 생성을 위해 확산 트랜스포머 모델을 탐구하기 위해 모든 코드와 모델을 https://github.com/whlzy/FiT 에 공개했습니다.

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.