확산 트랜스포머를 160억 파라미터로 확장하기

Scaling Diffusion Transformers to 16 Billion Parameters

July 16, 2024

저자: Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, Junshi Huang

cs.AI

초록

본 논문에서는 확산 트랜스포머의 희소 버전인 DiT-MoE를 소개하며, 이는 고밀도 네트워크와 경쟁력이 있으면서도 고도로 최적화된 추론을 보여줍니다. DiT-MoE는 두 가지 간단한 설계를 포함합니다: 공유 전문가 라우팅과 전문가 수준의 균형 손실로, 이를 통해 공통 지식을 포착하고 다양한 라우팅된 전문가 간의 중복성을 줄입니다. 조건부 이미지 생성에 적용할 때, 전문가 전문화에 대한 심층 분석은 몇 가지 흥미로운 관찰을 제공합니다: (i) 전문가 선택은 공간 위치와 노이즈 제거 시간 단계에 대한 선호도를 보이지만, 다른 클래스 조건 정보에는 둔감합니다; (ii) MoE 층이 깊어질수록, 전문가 선택은 특정 공간 위치에서 점점 분산과 균형으로 이동합니다. (iii) 전문가 전문화는 초기 시간 단계에서 더 집중되는 경향이 있으며, 이후 절반 이후에는 점점 균일해집니다. 우리는 이를 확산 과정에 기인한다고 보며, 이 과정은 먼저 저주파 공간 정보를 모델링한 후 고주파 복잡 정보를 모델링합니다. 위의 지침을 바탕으로, 일련의 DiT-MoE는 실험적으로 고밀도 네트워크와 동등한 성능을 달성하면서도 추론 중에 훨씬 적은 계산 부하를 요구합니다. 더욱 고무적으로, 우리는 합성된 이미지 데이터를 통해 DiT-MoE의 잠재력을 입증하며, 16.5B 매개변수로 확산 모델을 확장하여 512x512 해상도 설정에서 1.80의 새로운 SoTA FID-50K 점수를 달성했습니다. 프로젝트 페이지: https://github.com/feizc/DiT-MoE.

English

In this paper, we present DiT-MoE, a sparse version of the diffusion Transformer, that is scalable and competitive with dense networks while exhibiting highly optimized inference. The DiT-MoE includes two simple designs: shared expert routing and expert-level balance loss, thereby capturing common knowledge and reducing redundancy among the different routed experts. When applied to conditional image generation, a deep analysis of experts specialization gains some interesting observations: (i) Expert selection shows preference with spatial position and denoising time step, while insensitive with different class-conditional information; (ii) As the MoE layers go deeper, the selection of experts gradually shifts from specific spacial position to dispersion and balance. (iii) Expert specialization tends to be more concentrated at the early time step and then gradually uniform after half. We attribute it to the diffusion process that first models the low-frequency spatial information and then high-frequency complex information. Based on the above guidance, a series of DiT-MoE experimentally achieves performance on par with dense networks yet requires much less computational load during inference. More encouragingly, we demonstrate the potential of DiT-MoE with synthesized image data, scaling diffusion model at a 16.5B parameter that attains a new SoTA FID-50K score of 1.80 in 512times512 resolution settings. The project page: https://github.com/feizc/DiT-MoE.