让扩散模型的专家路由更关注关键细节,提升生成质量。
Focusing on What Matters: Saliency-Harnessing Accurate Routing for Diffusion MoE

- 用干净潜在特征指导路由,避开噪声干扰
- 多阶段损失约束计算分配,实现精准资源调度
- 无需重训练,适配已有扩散MoE模型
Mixture-of-Experts(MoE)架构已成为扩展视觉生成扩散模型的强大范式。近期工作聚焦于自适应分配不同令牌的计算资源以提升效率与性能。然而,我们发现现有扩散MoE框架存在路由分配问题:路由器无法准确将更多计算资源分配给显著性高的令牌。分析表明,该问题源于路由器在整个去噪过程中依赖被噪声污染的潜在特征,随机噪声掩盖了关键的结构与纹理信息,导致路由器难以有效区分显著性令牌。为此,我们提出SharpMoE,一种后训练框架,其具备基于显著性的精准路由机制,利用干净潜在特征作为无噪声引导信号进行路由。通过绕过噪声扭曲的输入,SharpMoE为路由器提供清晰的显著性指引,使在高噪声阶段也能识别显著性令牌。此外,我们引入轨迹路由损失,约束多步去噪轨迹中的计算分配,确保生成过程中的资源分配精确。大量实验表明,SharpMoE是一种通用、即插即用的解决方案,可进一步提升预训练完成的MoE模型,在视觉生成任务中达到当前最优表现。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。