O Roteamento Importa nos MoE: Dimensionando Transformadores de Difusão com Orientação Explícita de Roteamento

Resumo

A Mistura de Especialistas (MoE) emergiu como um paradigma poderoso para escalar a capacidade do modelo, preservando a eficiência computacional. Apesar do seu notável sucesso em grandes modelos de linguagem (LLMs), as tentativas existentes de aplicar MoE a Transformadores de Difusão (DiTs) têm produzido ganhos limitados. Atribuímos esta lacuna a diferenças fundamentais entre tokens linguísticos e visuais. Os tokens linguísticos são semanticamente densos, com variação pronunciada entre tokens, enquanto os tokens visuais exibem redundância espacial e heterogeneidade funcional, dificultando a especialização de especialistas em MoE visual. Para tal, apresentamos o ProMoE, uma estrutura MoE que apresenta um roteador de duas etapas com orientação de roteamento explícita que promove a especialização de especialistas. Especificamente, esta orientação incentiva o roteador a particionar os tokens de imagem em conjuntos condicionais e incondicionais via *routing* condicional de acordo com as suas funções, e a refinar as atribuições dos tokens de imagem condicionais através de *routing* prototípico com protótipos treináveis baseados no conteúdo semântico. Além disso, a alocação de especialistas baseada em similaridade no espaço latente, possibilitada pelo *routing* prototípico, oferece um mecanismo natural para incorporar orientação semântica explícita, e validamos que tal orientação é crucial para MoE visual. Com base nisto, propomos uma perda de contraste de roteamento que melhora explicitamente o processo de *routing* prototípico, promovendo coerência intra-especialista e diversidade inter-especialista. Extensos experimentos no *benchmark* ImageNet demonstram que o ProMoE supera os métodos mais avançados sob os objetivos de treino de Fluxo Retificado e DDPM. O código e os modelos serão disponibilizados publicamente.

English

Mixture-of-Experts (MoE) has emerged as a powerful paradigm for scaling model capacity while preserving computational efficiency. Despite its notable success in large language models (LLMs), existing attempts to apply MoE to Diffusion Transformers (DiTs) have yielded limited gains. We attribute this gap to fundamental differences between language and visual tokens. Language tokens are semantically dense with pronounced inter-token variation, while visual tokens exhibit spatial redundancy and functional heterogeneity, hindering expert specialization in vision MoE. To this end, we present ProMoE, an MoE framework featuring a two-step router with explicit routing guidance that promotes expert specialization. Specifically, this guidance encourages the router to partition image tokens into conditional and unconditional sets via conditional routing according to their functional roles, and refine the assignments of conditional image tokens through prototypical routing with learnable prototypes based on semantic content. Moreover, the similarity-based expert allocation in latent space enabled by prototypical routing offers a natural mechanism for incorporating explicit semantic guidance, and we validate that such guidance is crucial for vision MoE. Building on this, we propose a routing contrastive loss that explicitly enhances the prototypical routing process, promoting intra-expert coherence and inter-expert diversity. Extensive experiments on ImageNet benchmark demonstrate that ProMoE surpasses state-of-the-art methods under both Rectified Flow and DDPM training objectives. Code and models will be made publicly available.

O Roteamento Importa nos MoE: Dimensionando Transformadores de Difusão com Orientação Explícita de Roteamento

Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance

Resumo

Support