Сосредоточение на главном: точная маршрутизация на основе салиентности для диффузионной MoE

Аннотация

Архитектуры смеси экспертов (Mixture-of-Experts, MoE) стали мощной парадигмой для масштабирования диффузионных моделей в визуальной генерации. Последние достижения были сосредоточены на адаптивном распределении вычислительных ресурсов между разнородными токенами для повышения эффективности и производительности. Однако мы выявляем проблему назначения маршрутизации в существующих диффузионных MoE-фреймворках: маршрутизатор не способен точно выделить больше вычислительных ресурсов значимым токенам. Наш анализ связывает эту неудачу с зависимостью маршрутизатора от зашумленных латентных признаков на протяжении всего процесса шумоподавления. Такой стохастический шум скрывает критическую структурную и текстурную информацию, тем самым не позволяя маршрутизатору эффективно различать значимые токены. Для решения этой проблемы мы предлагаем SharpMoE — фреймворк после обучения с механизмом точной маршрутизации, использующим значимость, который применяет чистые латентные признаки в качестве безшумного направляющего сигнала для маршрутизации. Обходя искаженные шумом входные данные, SharpMoE предоставляет маршрутизатору четкое руководство по значимости, позволяя идентифицировать значимые токены даже на этапах с высоким уровнем шума. Кроме того, мы вводим потерю маршрутизации траектории, чтобы ограничить распределение вычислений на протяжении многошаговой траектории шумоподавления, обеспечивая точное выделение ресурсов по ходу генерации. Обширные эксперименты показывают, что SharpMoE служит универсальным решением типа "включай и работай", которое дополнительно улучшает предварительно обученные и сходящиеся MoE-модели, достигая передовых результатов в визуальной генерации.

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.