Se concentrer sur l'essentiel : Routage précis exploitant la saillance pour le MoE de diffusion

Résumé

Les architectures de Mélange d'Experts (MoE) sont apparues comme un paradigme puissant pour le passage à l'échelle des modèles de diffusion en génération visuelle. Les avancées récentes se sont concentrées sur l'allocation adaptative des ressources de calcul entre différents tokens afin d'améliorer l'efficacité et les performances. Cependant, nous identifions un problème d'affectation de routage dans les cadres MoE de diffusion existants : le routeur ne parvient pas à allouer précisément davantage de ressources de calcul aux tokens saillants. Notre analyse attribue cet échec au fait que le routeur s'appuie sur des caractéristiques latentes corrompues par le bruit tout au long du processus de débruitage. Ce bruit stochastique obscurcit les informations structurelles et texturales critiques, empêchant ainsi le routeur de distinguer efficacement les tokens saillants. Pour y remédier, nous proposons SharpMoE, un cadre post-entraînement doté d'un mécanisme de routage précis exploitant la saillance, qui utilise des caractéristiques latentes propres comme signal de guidage sans bruit pour le routage. En contournant les entrées déformées par le bruit, SharpMoE fournit au routeur un guidage clair en termes de saillance, permettant l'identification des tokens saillants même dans les étapes à fort bruit. De plus, nous introduisons une perte de routage de trajectoire pour contraindre l'allocation de calcul tout au long de la trajectoire de débruitage en plusieurs étapes, garantissant une répartition précise des ressources lors du déroulement de la génération. Des expériences approfondies démontrent que SharpMoE constitue une solution polyvalente et prête à l'emploi qui améliore encore les modèles MoE pré-entraînés et convergés, atteignant des performances de pointe en génération visuelle.

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.