Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving Routing in Sparse Mixture of Experts with Graph of Tokens

Tam Nguyen, Ngoc N. Tran, Khai Nguyen, Richard G. Baraniuk

TL;DR

The paper identifies routing fluctuations in Sparse MoE as a consequence of token-wise independence in expert selection. It introduces Similarity-Aware (S)MoE and Attention-Aware (S)MoE, leveraging token similarities and attention-derived dependencies to reduce routing entropy and stabilize expert assignments. Through a Probabilistic Graphical Model lens and extensions to MoE-Transformer blocks, the authors prove theoretical entropy reductions and demonstrate empirical gains across vision, language, and downstream tasks, including robustness to perturbations. The proposed methods are modular, improving routing stability and load balancing while preserving or enhancing performance, making them practical plug-and-play enhancements for contemporary MoE architectures.

Abstract

Sparse Mixture of Experts (SMoE) has emerged as a key to achieving unprecedented scalability in deep learning. By activating only a small subset of parameters per sample, SMoE achieves an exponential increase in parameter counts while maintaining a constant computational overhead. However, SMoE models are susceptible to routing fluctuations--changes in the routing of a given input to its target expert--at the late stage of model training, leading to model non-robustness. In this work, we unveil the limitation of SMoE through the perspective of the probabilistic graphical model (PGM). Through this PGM framework, we highlight the independence in the expert-selection of tokens, which exposes the model to routing fluctuation and non-robustness. Alleviating this independence, we propose the novel Similarity-Aware (S)MoE, which considers interactions between tokens during expert selection. We then derive a new PGM underlying an (S)MoE-Attention block, going beyond just a single (S)MoE layer. Leveraging the token similarities captured by the attention matrix, we propose the innovative Attention-Aware (S)MoE, which employs the attention matrix to guide the routing of tokens to appropriate experts in (S)MoE. We theoretically prove that Similarity/Attention-Aware routing help reduce the entropy of expert selection, resulting in more stable token routing mechanisms. We empirically validate our models on various tasks and domains, showing significant improvements in reducing routing fluctuations, enhancing accuracy, and increasing model robustness over the baseline MoE-Transformer with token routing via softmax gating.

Improving Routing in Sparse Mixture of Experts with Graph of Tokens

TL;DR

The paper identifies routing fluctuations in Sparse MoE as a consequence of token-wise independence in expert selection. It introduces Similarity-Aware (S)MoE and Attention-Aware (S)MoE, leveraging token similarities and attention-derived dependencies to reduce routing entropy and stabilize expert assignments. Through a Probabilistic Graphical Model lens and extensions to MoE-Transformer blocks, the authors prove theoretical entropy reductions and demonstrate empirical gains across vision, language, and downstream tasks, including robustness to perturbations. The proposed methods are modular, improving routing stability and load balancing while preserving or enhancing performance, making them practical plug-and-play enhancements for contemporary MoE architectures.

Abstract

Sparse Mixture of Experts (SMoE) has emerged as a key to achieving unprecedented scalability in deep learning. By activating only a small subset of parameters per sample, SMoE achieves an exponential increase in parameter counts while maintaining a constant computational overhead. However, SMoE models are susceptible to routing fluctuations--changes in the routing of a given input to its target expert--at the late stage of model training, leading to model non-robustness. In this work, we unveil the limitation of SMoE through the perspective of the probabilistic graphical model (PGM). Through this PGM framework, we highlight the independence in the expert-selection of tokens, which exposes the model to routing fluctuation and non-robustness. Alleviating this independence, we propose the novel Similarity-Aware (S)MoE, which considers interactions between tokens during expert selection. We then derive a new PGM underlying an (S)MoE-Attention block, going beyond just a single (S)MoE layer. Leveraging the token similarities captured by the attention matrix, we propose the innovative Attention-Aware (S)MoE, which employs the attention matrix to guide the routing of tokens to appropriate experts in (S)MoE. We theoretically prove that Similarity/Attention-Aware routing help reduce the entropy of expert selection, resulting in more stable token routing mechanisms. We empirically validate our models on various tasks and domains, showing significant improvements in reducing routing fluctuations, enhancing accuracy, and increasing model robustness over the baseline MoE-Transformer with token routing via softmax gating.
Paper Structure (38 sections, 3 theorems, 38 equations, 5 figures, 10 tables)

This paper contains 38 sections, 3 theorems, 38 equations, 5 figures, 10 tables.

Table of Contents

  1. Introduction
  2. Sparse Mixture of Experts
  3. Contributions
  4. Routing in SMoE with Graph of Tokens
  5. A Probabilistic Graphical Model for (S)MoE
  6. Similarity-Aware (S)MoE
  7. Beyond One Layer: Routing in MoE Transformer with Graph of Tokens
  8. PGM for MoE Transformer
  9. Multi-head Attention and MoE
  10. Graphical model for (S)MoE-Attention
  11. Attention-Aware (S)MoE)
  12. Entropy Analysis of Similarity and Attention-Aware Routing
  13. Experimental Results
  14. Related Work
  15. Concluding Remarks
  16. ...and 23 more sections

Key Result

Lemma 1

The distribution of the expert selection $e^a_i$ conditioned on ${\mathbf U}, {\mathbf X}$, is given by where the posteriors with the prior $\bold{A}_{h}[i, j] := {\mathbb P}(z_i = j | h_i = h, {\mathbf X})$ and $\bold{H}[i, h] := {\mathbb P}(h_i = h)$ and the likelihood $\mathbf{L}_{h}[i, j] := \mathcal{N}({\bm u}_i |{\mathbf W}_h{\bm x}_{j}, \sigma^2\bold{I})$. This results in

Figures (5)

  • Figure 1: PGMs for (S)MoE ($\mathcal{G}_1$) and Similarity-Aware (S)MoE ($\mathcal{G}_2$). Directed paths are shown by arrows; dotted arrow indicates concatenation; blue arrows highlight differences. ${\mathbf U} = [{\bm u}_1, \dots, {\bm u}_N]^T$ is the input sequence of (S)MoE and Similarity-Aware (S)MoE. ${\bm u}_i, {\bm o}_i$ are input-output at position $i \in [1,...,N]$. Variables $e_i$ and $e^s_i$ denote expert selection for ${\bm u}_i$ in (S)MoE and Similarity-Aware (S)MoE, respectively; $s_i$ represents similarity variable for ${\mathbf U}$.
  • Figure 2: PGMs for (S)MoE-Attention ($\mathcal{G}_3$) and Attention-Aware (S)MoE ($\mathcal{G}_4$) defined in Def. \ref{['def:mam_pgm']} and Def. \ref{['def:aamm_pgm']}, respectively. Directed paths shown by arrows; dotted arrow indicates concatenation; blue arrows highlight differences.
  • Figure 3: Comparison of routing fluctuation and entropy ratio across layers for Baseline SMoE, Attention-Aware SMoE, and Similarity-Aware SMoE
  • Figure 4: Comparison of expert routing distribution for baseline SMoE, Attention-Aware SMoE, and Similarity-Aware SMoE
  • Figure 5: Comparison of Routing Fluctuation and Entropy Ratio Across Layers for Baseline SmoE Top-1, Attention-Aware SMoE Top-1, and Similarity-Aware SMoE Top-1

Theorems & Definitions (11)

  • Remark 1
  • Remark 2: Deriving Other Routers
  • Definition 1
  • Remark 3
  • Definition 2
  • Definition 3
  • Definition 4: Attention-Aware MoE Generative Model (A$^{2}$MM)
  • Lemma 1
  • Definition 5: Attention-Aware (S)MoE.
  • Proposition 1
  • ...and 1 more