SwiftBrush v2: 당신의 One-step 확산 모델을 그것의 교사보다 더 나아지게 만들다

SwiftBrush v2: Make Your One-step Diffusion Model Better Than Its Teacher

August 26, 2024

저자: Trung Dao, Thuan Hoang Nguyen, Thanh Le, Duc Vu, Khoi Nguyen, Cuong Pham, Anh Tran

cs.AI

초록

본 논문에서는 두 단계의 안정적 확산 모델인 SwiftBrush의 성능을 향상시켜 다단계 안정적 확산 모델과 경쟁력을 갖도록 목표로 합니다. 우선, SwiftBrush와 SD Turbo 간의 품질-다양성 트레이드오프를 탐구합니다: 전자는 이미지 다양성에서 뛰어나지만 후자는 이미지 품질에서 뛰어납니다. 이 관찰은 더 나은 가중치 초기화와 효율적인 LoRA 훈련을 포함한 훈련 방법론의 수정을 제안하게 했습니다. 게다가, 새로운 클램프드 CLIP 손실을 도입하여 이미지-텍스트 정렬을 향상시키고 이미지 품질을 향상시킵니다. 놀랍게도, 효율적인 LoRA 및 전체 훈련으로 훈련된 모델의 가중치를 결합함으로써 새로운 최첨단 단계의 확산 모델을 달성하며 FID가 8.14로 모든 GAN 기반 및 다단계 안정적 확산 모델을 능가합니다. 평가 코드는 다음에서 확인할 수 있습니다: https://github.com/vinairesearch/swiftbrushv2.

English

In this paper, we aim to enhance the performance of SwiftBrush, a prominent one-step text-to-image diffusion model, to be competitive with its multi-step Stable Diffusion counterpart. Initially, we explore the quality-diversity trade-off between SwiftBrush and SD Turbo: the former excels in image diversity, while the latter excels in image quality. This observation motivates our proposed modifications in the training methodology, including better weight initialization and efficient LoRA training. Moreover, our introduction of a novel clamped CLIP loss enhances image-text alignment and results in improved image quality. Remarkably, by combining the weights of models trained with efficient LoRA and full training, we achieve a new state-of-the-art one-step diffusion model, achieving an FID of 8.14 and surpassing all GAN-based and multi-step Stable Diffusion models. The evaluation code is available at: https://github.com/vinairesearch/swiftbrushv2.