Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Selective Sinkhorn Routing for Improved Sparse Mixture of Experts

Duc Anh Nguyen, Huu Binh Ta, Nhuan Le Duc, Tan M. Nguyen, Toan Tran

TL;DR

The paper introduces Selective Sinkhorn Routing (SSR) to improve Sparse Mixture of Experts by formulating token-to-expert routing as entropy-regularized optimal transport with balance constraints. SSR derives routing weights from the OT transport map and applies Sinkhorn updates sparsely during training, avoiding reliance on auxiliary balancing losses. The approach yields faster convergence, higher accuracy, and greater robustness in both language modeling and image classification benchmarks, while preserving flexibility in expert selection. Theoretical analysis and extensive experiments demonstrate that limited Sinkhorn routing, combined with optional noise for exploration, improves training stability and balance without sacrificing inference determinism.

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.

Selective Sinkhorn Routing for Improved Sparse Mixture of Experts

TL;DR

The paper introduces Selective Sinkhorn Routing (SSR) to improve Sparse Mixture of Experts by formulating token-to-expert routing as entropy-regularized optimal transport with balance constraints. SSR derives routing weights from the OT transport map and applies Sinkhorn updates sparsely during training, avoiding reliance on auxiliary balancing losses. The approach yields faster convergence, higher accuracy, and greater robustness in both language modeling and image classification benchmarks, while preserving flexibility in expert selection. Theoretical analysis and extensive experiments demonstrate that limited Sinkhorn routing, combined with optional noise for exploration, improves training stability and balance without sacrificing inference determinism.

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.

Paper Structure

This paper contains 34 sections, 6 theorems, 40 equations, 6 figures, 11 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Sparse Mixture of Experts
  4. Routing in Sparse Mixture of Experts
  5. Preliminaries
  6. Sparse Mixture of Experts
  7. Sinkhorn-Knopp Algorithm
  8. Method
  9. Token-to-expert Assignment as an Entropy-regularized Optimal Transport Problem
  10. Selective Sinkhorn Routing for Sparse Mixture of Experts
  11. Experiments
  12. Settings
  13. Datasets
  14. Baselines
  15. Evaluation Metrics and Analysis
  16. ...and 19 more sections

Key Result

Proposition 4.1

Let $\boldsymbol{\hat{\Pi}} \in \mathbb{R}^{m \times n}$ be the solution to the entropy-regularized optimal transport problem in Eq. OT_lb. For each token $i \in \{1, \ldots, m\}$, suppose we are allowed to assign it to at most $k$ experts, with routing weights ${\alpha}_i \in \mathbb{R}^n$ satisfyi We select the top-$k$ experts $E_{i_1}, E_{i_2}, \ldots, E_{i_k}$ with the highest transport scores

Figures (6)

  • Figure 1: Training of SSR. At each layer during the forward pass, a fixed probability $p \in [0,1]$ is introduced. We randomly choose Sinkhorn-based routing with probability $p$ and Softmax-based routing with probability $1-p$.
  • Figure 2: Training time Overhead vs. SMoE (Vanilla) over 1-Head setting on Wikitext-103 dataset.
  • Figure 3: Performance of SSR-S on Enwik-8 under varying $\xi$ and $\alpha_{\text{noise}}$ values.
  • Figure 4: Training time Overhead vs. Vanilla SMoE over 2-Head setting on Wikitext-103 dataset.
  • Figure 5: Training time Overhead vs. Vanilla SMoE over 4-Head setting on Wikitext-103 dataset.
  • ...and 1 more figures

Theorems & Definitions (9)

  • Proposition 4.1
  • Proposition 4.2: Load Balancing Is Valid During Training but Not Inference
  • Proposition 4.3: Noise Ensures Every Expert Has Nonzero Selection Probability
  • Proposition 5.1
  • proof
  • Proposition 5.2
  • proof
  • Proposition 5.3
  • proof