Focando no que Importa: Roteamento Preciso com Aproveitamento de Saliência para MoE de Difusão

Resumo

Arquiteturas de Mistura de Especialistas (MoE) emergiram como um paradigma poderoso para escalar modelos de difusão na geração visual. Avanços recentes têm se concentrado na alocação adaptativa de recursos computacionais entre tokens diversos para melhorar a eficiência e o desempenho. No entanto, identificamos um problema de atribuição de roteamento em estruturas MoE de difusão existentes: o roteador falha em alocar com precisão mais recursos computacionais para tokens salientes. Nossa análise atribui essa falha à dependência do roteador em características latentes corrompidas por ruído ao longo do processo de remoção de ruído. Esse ruído estocástico obscurece as informações estruturais e texturais críticas, impedindo assim que o roteador distinga efetivamente os tokens salientes. Para resolver isso, propomos o SharpMoE, uma estrutura pós-treinamento com um mecanismo de roteamento preciso que aproveita a saliência, utilizando características latentes limpas como um sinal de orientação livre de ruído para o roteamento. Ao contornar as entradas distorcidas por ruído, o SharpMoE fornece ao roteador uma orientação clara de saliência, permitindo a identificação de tokens salientes mesmo em estágios de alto ruído. Além disso, introduzimos uma perda de roteamento de trajetória para restringir a alocação de computação ao longo da trajetória de remoção de ruído em várias etapas, garantindo uma alocação precisa de recursos ao longo da implementação da geração. Experimentos extensivos demonstram que o SharpMoE atua como uma solução versátil e plug-and-play que aprimora ainda mais os modelos MoE pré-treinados e convergidos, alcançando desempenho de estado da arte na geração 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.