중요한 것에 집중하기: 확산 MoE를 위한 현저성 활용 정확한 라우팅

초록

Mixture-of-Experts (MoE) 아키텍처는 시각 생성에서 확산 모델을 확장하기 위한 강력한 패러다임으로 부상했다. 최근의 발전은 효율성과 성능을 개선하기 위해 다양한 토큰에 걸쳐 계산 자원을 적응적으로 할당하는 데 초점을 맞추고 있다. 그러나 우리는 기존 확산 MoE 프레임워크에서 라우팅 할당 문제를 식별한다: 라우터가 중요한 토큰에 더 많은 계산 자원을 정확히 할당하지 못한다. 우리의 분석은 이 실패를 라우터가 노이즈 제거 과정 전반에 걸쳐 노이즈로 오염된 잠재 특징에 의존하기 때문이라고 설명한다. 이러한 확률적 노이즈는 중요한 구조적 및 질감 정보를 모호하게 하여, 라우터가 중요한 토큰을 효과적으로 구별하지 못하게 한다. 이를 해결하기 위해, 우리는 SharpMoE를 제안한다. 이는 중요도를 활용한 정확한 라우팅 메커니즘을 갖춘 사후 훈련 프레임워크로, 깨끗한 잠재 특징을 노이즈가 없는 라우팅 안내 신호로 활용한다. 노이즈로 왜곡된 입력을 우회함으로써, SharpMoE는 라우터에 명확한 중요도 안내를 제공하여 높은 노이즈 단계에서도 중요한 토큰을 식별할 수 있게 한다. 또한, 우리는 다단계 노이즈 제거 궤적 전반에 걸쳐 계산 할당을 제약하는 궤적 라우팅 손실(trajectory routing loss)을 도입하여, 생성 롤아웃 과정에서 정밀한 자원 할당을 보장한다. 광범위한 실험을 통해 SharpMoE가 사전 훈련되고 수렴된 MoE 모델을 더욱 향상시키는 다재다능한 플러그 앤 플레이 솔루션으로 작용하여, 시각 생성에서 최첨단 성능을 달성함을 입증한다.

English

Mixture-of-Experts (MoE) architectures have emerged as a powerful paradigm for scaling diffusion models in visual generation. Recent advancements have focused on adaptively allocating computational resources across diverse tokens to improve efficiency and performance. However, we identify a routing assignment problem in existing diffusion MoE frameworks: the router fails to accurately allocate more computational resources to salient tokens. Our analysis attributes this failure to the router's reliance on noise-corrupted latent features throughout the denoising process. Such stochastic noise obscures the critical structural and textural information, thereby preventing the router from effectively distinguishing salient tokens. To address this, we propose SharpMoE, a post-training framework with a saliency-harnessing accurate routing mechanism, which utilizes clean latent features as a noise-free guidance signal for routing. By bypassing the noise-distorted inputs, SharpMoE provides the router with clear saliency guidance, enabling the identification of salient tokens even in high-noise stages. Furthermore, we introduce a trajectory routing loss to constrain the compute allocation throughout the multi-step denoising trajectory, ensuring precise resource allocation along the generation rollout. Extensive experiments demonstrate that SharpMoE serves as a versatile, plug-and-play solution that further enhances the pretrained, converged MoE models, achieving state-of-the-art performance in visual generation.