Selective Sinkhorn Routing for Improved Sparse Mixture of Experts

TL;DR

This paper introduces Selective Sinkhorn Routing, a novel method for sparse mixture of experts that improves expert balancing and model performance by formulating token assignment as an optimal transport problem, reducing training overhead.

Contribution

The paper proposes SSR, a Sinkhorn-based routing mechanism that replaces auxiliary losses, enabling more effective expert balancing with less complexity and faster training.

Findings

SSR achieves higher accuracy in language and image tasks.

SSR reduces training time compared to previous methods.

SSR enhances robustness to input corruption.

Abstract

Sparse Mixture-of-Experts (SMoE) has gained prominence as a scalable and computationally efficient architecture, enabling significant growth in model capacity without incurring additional inference costs. However, existing SMoE models often rely on auxiliary losses (e.g., z-loss, load balancing) and additional trainable parameters (e.g., noisy gating) to encourage expert diversity, leading to objective misalignment and increased model complexity. Moreover, existing Sinkhorn-based methods suffer from significant training overhead due to their heavy reliance on the computationally expensive Sinkhorn algorithm. In this work, we formulate token-to-expert assignment as an optimal transport problem, incorporating constraints to ensure balanced expert utilization. We demonstrate that introducing a minimal degree of optimal transport-based routing enhances SMoE performance without requiring auxiliary balancing losses. Unlike previous methods, our approach derives gating scores directly from the transport map, enabling more effective token-to-expert balancing, supported by both theoretical analysis and empirical results. Building on these insights, we propose Selective Sinkhorn Routing (SSR), a routing mechanism that replaces auxiliary loss with lightweight Sinkhorn-based routing. SSR promotes balanced token assignments while preserving flexibility in expert selection. Across both language modeling and image classification tasks, SSR achieves faster training, higher accuracy, and greater robustness to input corruption.