Влияние маршрутизации в MoE: масштабирование диффузионных трансформаторов с явным управлением маршрутизацией

Аннотация

Смесь экспертов (MoE) стала мощной парадигмой для масштабирования емкости моделей при сохранении вычислительной эффективности. Несмотря на заметный успех в больших языковых моделях (LLM), существующие попытки применения MoE к диффузионным трансформерам (DiT) принесли ограниченный выигрыш. Мы объясняем этот разрыв фундаментальными различиями между языковыми и визуальными токенами. Языковые токены семантически плотны с выраженной вариативностью между токенами, тогда как визуальные токены демонстрируют пространственную избыточность и функциональную гетерогенность, что затрудняет специализацию экспертов в визуальном MoE. В связи с этим мы представляем ProMoE — фреймворк MoE с двухэтапным маршрутизатором и явными рекомендациями по маршрутизации, которые способствуют специализации экспертов. В частности, эти рекомендации побуждают маршрутизатор разделять изображения на условные и безусловные множества посредством условной маршрутизации в соответствии с их функциональными ролями, а также уточнять назначения условных токенов изображения через прототипную маршрутизацию с обучаемыми прототипами на основе семантического содержания. Более того, распределение экспертов на основе сходства в латентном пространстве, обеспечиваемое прототипной маршрутизацией, предлагает естественный механизм для включения явной семантической guidance, и мы подтверждаем, что такая guidance критически важна для визуального MoE. На основе этого мы предлагаем контрастную потерю для маршрутизации, которая явно усиливает процесс прототипной маршрутизации, способствуя внутриэкспертной когерентности и межэкспертному разнообразию. Обширные эксперименты на benchmark ImageNet демонстрируют, что ProMoE превосходит современные методы как при использовании Rectified Flow, так и DDPM в качестве целей обучения. Код и модели будут общедоступны.

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.

Влияние маршрутизации в MoE: масштабирование диффузионных трансформаторов с явным управлением маршрутизацией

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

Аннотация

Support