Centrándose en lo que importa: enrutamiento preciso mediante el aprovechamiento de la saliencia para MoE de difusión

Resumen

Las arquitecturas de Mezcla de Expertos (MoE) han surgido como un paradigma poderoso para escalar modelos de difusión en la generación visual. Avances recientes se han centrado en asignar adaptativamente recursos computacionales entre diversos tokens para mejorar la eficiencia y el rendimiento. Sin embargo, identificamos un problema de asignación de enrutamiento en los marcos MoE de difusión existentes: el enrutador no logra asignar con precisión más recursos computacionales a los tokens salientes. Nuestro análisis atribuye este fracaso a la dependencia del enrutador de características latentes contaminadas por ruido durante todo el proceso de eliminación de ruido. Dicho ruido estocástico oscurece la información estructural y textural crítica, impidiendo así que el enrutador distinga eficazmente los tokens salientes. Para abordar esto, proponemos SharpMoE, un marco de post-entrenamiento con un mecanismo de enrutamiento preciso que aprovecha la prominencia, el cual utiliza características latentes limpias como una señal de guía libre de ruido para el enrutamiento. Al evitar las entradas distorsionadas por el ruido, SharpMoE proporciona al enrutador una guía de prominencia clara, permitiendo la identificación de tokens salientes incluso en etapas de alto ruido. Además, introducimos una pérdida de enrutamiento de trayectoria para restringir la asignación computacional a lo largo de la trayectoria de eliminación de ruido de múltiples pasos, asegurando una asignación precisa de recursos durante el desarrollo de la generación. Experimentos exhaustivos demuestran que SharpMoE sirve como una solución versátil y plug-and-play que mejora aún más los modelos MoE preentrenados y convergentes, logrando un rendimiento de vanguardia en la generación visual.

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.