Mutual-Inform SMoE: Improving Routing Stability via Probabilistic Graphical Model

Authors: Tam Minh Nguyen, Ngoc N. Tran, Khai Nguyen, Richard Baraniuk
Primary area: foundation or frontier models, including LLMs
Keywords: mixture of expert, transformer, probabilistic graphical model, robustness

Abstract:
Sparse Mixture of Experts (SMoE) has emerged as a breakthrough approach for achieving unprecedented scalability in deep learning. By enabling models to expand their parameter count exponentially while selectively activating only a small subset of parameters per sample, SMoEs maintain high efficiency. However, SMoE models are susceptible to routing fluctuations, leading to instability and non-robustness. In this work, we unveils SMoE-based attention as a point estimate of a regression function of a 3-layer hierarchical mixture of experts regression. Through this probabilistic graphical model (PGM) framework, we highlight the conditional independence in expert-selection process of tokens, which exposes the model to routing fluctuation and non-robustness. Motivating by this PGM framework, we propose Mutual-Inform SMoEs, including Similarity and Attention-Inform SMoE, which eliminate the assumption of conditional independence by allowing tokens to directly influence each other on expert-decisions. We theoretically demonstrate that our methods lower the entropy in decision-making, enabling more confident and consistent expert assignments. Finally, we empirically validate our models on ImageNet classification and Wikitext-103 language modeling, showing significant improvements in reducing routing fluctuations, enhancing performance, and increasing model robustness compared to baseline Transformer-SMoE models.

Decision: reject
Reviewer scores: 5, 5, 5, 8

Reviews and discussion:
<thread>
creation date:
1730579794697

thread_type:
official_review

<thread_object>
creator: reviewer_azaz

<content>
summary:
This work addresses a core limitation of SMoEs—routing instability—by introducing collaborative token routing mechanisms, enhancing both model robustness and efficiency. This is show both theoretically and in a few benchmarks.
	
**Contributions:**

Probabilistic Graphical Model (PGM) Framework: Introduces a PGM framework to model SMoE-based attention, showing that expert selection is conditionally independent for each token, leading to unstable routing.

Mutual-Inform SMoE Models: Proposes two new SMoE variants:
- Similarity-Inform SMoE: Routes similar tokens to the same expert, using token similarities from the MoE layer.
- Attention-Inform SMoE: Utilizes relationships from the attention layer to inform expert assignments, aligning expert choice with token interactions.

This is validated theoretically (with entropy reduction) and empirically (with task specific benchmarks)

strengths:
**Originality:** Introduces Mutual-Inform SMoE, a novel approach for stabilizing SMoE routing using token similarity and attention dependencies. Innovatively reinterprets SMoE through a probabilistic graphical model (PGM), which is a fresh way of addressing routing challenges.

**Quality:** Very rigorous theoretical framework (congrats!); sound mathematical proofs back claims about reduced entropy and improved stability. 

**Significance:** Addresses a critical issue in SMoEs (routing instability), with potential to impact scalable model design for large-scale LLM and vision applications. Improvements in both standard and adversarial settings underscore its value in robust model deployment.

weaknesses:
### Presentation weaknesses:
- Writing needs to be considerably improved (sentences unclear, leading to confusions)
- Many typos need fixing (eg: Abstract: In this work, we unveil**s**; Motivat**ing** by this PGM framework…)
- Proofs could be explained a bit better than just a series of latex equations.

### Empirical work weaknesses
- The empirical validation is weak, more benchmarks would be appreciated. The experiments could be strengthened by comparing against more recent or diverse baselines
- Mutual-Inform SMoE models show strong results, additional ablations would clarify the impact of individual components (e.g., temperature parameters, entropy reduction mechanisms). Including these would highlight the contributions of each module in performance gains.
- There is a lack of discussion on the computational costs of the method, and the computation penalty for using these more performant models compared to the baseline.

questions:
- How does this work compare to other, more modern methods and across more benchmarks?
- What is the computational cost of these methods?

(See weaknesses for details)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730643647619

thread_type:
official_review

<thread_object>
creator: reviewer_mg8s

<content>
summary:
This paper proposes a novel perspective on Sparse Mixture of Experts (SMoE) through a probabilistic graphical model (PGM) framework. The authors identify that the conditional independence in expert selection for tokens leads to routing fluctuations and model instability. To address this, they introduce Mutual-Inform SMoE with two variants: Similarity-Inform and Attention-Inform, which allow tokens to influence each other's expert assignments based on their similarities or attention patterns. They provide theoretical analysis showing their method reduces entropy in routing decisions and demonstrate empirical improvements on ImageNet classification and WikiText-103 language modeling tasks.

strengths:
1. The paper provides a novel theoretical foundation for understanding SMoE through PGM, offering valuable insights into the routing fluctuation problem and presenting a well-grounded solution. The mathematical derivations and proofs are rigorous and well-presented.

2. The proposed solution is elegant and practical, requiring minimal additional computational overhead while showing significant improvements in both performance and stability. The two variants (Similarity-Inform and Attention-Inform) provide flexibility in implementation.

3. The experimental evaluation is comprehensive, covering both vision and language tasks, with thorough analysis of routing stability, entropy reduction, and robustness against various types of perturbations.

weaknesses:
1. While the paper demonstrates improvements in routing stability, it doesn't fully explore the trade-offs between stability and adaptability. A more detailed analysis of whether increased stability might sometimes come at the cost of reduced model flexibility would be valuable.

2. The hyperparameter sensitivity analysis is limited, particularly regarding the temperature parameter $\tau$ in Similarity-Inform SMoE and $\sigma$ in Attention-Inform SMoE. Understanding how these parameters affect performance would be crucial for practical implementations.

3. The experiments focus primarily on medium-sized models. Given that routing issues often become more pronounced in larger-scale settings, evaluation on larger models would strengthen the paper's claims.

4. While the paper mentions load balancing benefits in the appendix, this aspect deserves more attention in the main text, as it could be a significant practical advantage of the proposed approach.

questions:
1. How does the approach scale with very large numbers of experts (e.g., hundreds or thousands)?

2. Have you explored the possibility of dynamically adjusting the influence of token similarities based on training progress?

3. How does the method perform when applied to multi-modal tasks where routing patterns might be more complex?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730756280867

thread_type:
official_review

<thread_object>
creator: reviewer_jpxc

<content>
summary:
The paper proposes "MUTUAL-INFORM SMOEs", MoEs with a tweaked routing mechanism that allows for information sharing between routed tokens. The authors motivate their proposed method as a solution to reducing the routing fluctuation problem in MoEs. They provide theoretical justification for their method's reduction in routing fluctuation by showing that the entropy of routing decisions of "MUTUAL-INFORM SMOEs" is upper bounded by standard topK + softmax routing. The authors provide additional experiments on language modeling (standard + adversarially perturbed text) and image classification, showing that their method outperforms a standard topK + softmax baseline in routing fluctuation and validation perplexity.

strengths:
- The Attention-Inform (S)MoE and Similarity-Inform (S)MoE routing functions are novel and potentially interesting to the community. The introduction of routing mechanisms that can share information between tokens can lead to better-performing MoEs.
- It is nice to see the theoretical justification for the proposed method.
- I like the idea of including an adversarial test set. 
- Providing code for the method is appreciated and helps validate the results.

weaknesses:
My **main concern** is the lack of comparison to previous work that also improves the *routing fluctuation* problem of MoEs. Authors cite [1,2] numerous times when referring to the routing fluctuation problem. However, they do not provide a comparison to Stablemoe or X-MoE. To the best of my knowledge, [1,2] are the only works that claim this *routing fluctuation* problem exists in MoEs and to provide a solution for it. As clear follow-up work to [1,2], both more than two years old, it is very important to directly compare to their method. Adding a comparison to these methods would greatly strengthen the contributions of the work. 

**Other concerns** 
- The method relies on the attention matrix for similarity scores between different tokens. How does this influence the computational and memory complexity of an MoE forward pass? Providing figures that show the timings with respect to a baseline would address this concern.
- The introduction is structured unconventionally and I found myself taking time to find the motivation for the work.
- It would be helpful to understand how the performance of M-I SMoE changes as the number of experts is increased, the granularity of the experts changes, and the number of active parameters changes. 
- The architecture of the MoEs used (e.g., number of experts per layer) should be reported in the main text, not only in the appendix.
- I am unable to find the dimensions of the transformer used in the language modeling experiments.


Typos:
- "Motivating by this PGM framework" Motivating --> Motivated
- Page 10: "We leave it for future work. We leave it for future work"


[1] Stablemoe: Stable routing strategy for mixture of experts ( https://aclanthology.org/2022.acl-long.489/ )
[2] On the Representation Collapse of Sparse Mixture of Experts ( https://arxiv.org/pdf/2204.09179 )

questions:
- Is routing fluctuation discussed as a problem for MoEs outside of papers [1,2]?


[1] Stablemoe: Stable routing strategy for mixture of experts ( https://aclanthology.org/2022.acl-long.489/ )
[2] On the Representation Collapse of Sparse Mixture of Experts ( https://arxiv.org/pdf/2204.09179 )
</content>
</thread_object>

</thread>
<thread>
creation date:
1731141016721

thread_type:
official_review

<thread_object>
creator: reviewer_rswg

<content>
summary:
This paper formalized the transformer attention head as a probabilistic graph model and suggested that independent expert selections for different tokens in Sparse MoE contributes to the routing fluctuations. Based on such analyses, the authors proposed two intuitive approaches of aggregating other tokens' expert selections into the decision of the current token with either similarity between output features or adapted attention scores as the weights. Experiments in the paper demonstrated that such techniques could bring the improvements in terms of model performance and routing stability.

strengths:
The structure of the paper is well-organized and comprehensive. The paper clearly identified the potential issues behind routing fluctuations at start and formed two solutions step by step.

weaknesses:
1. The paper's writing did not fully meet the standard of a professional research publication . 
- Lack of rigorous and concise formula definitions although the authors tried to analyze the problems in a formal way. 
**a**. Diverse meanings of the same subscript in different formula. In Line 75-77, the subscript should be $k$ rather than $i$. It is not appropriate to use $i$ in line 87. Please do not confuse these two index symbols and try to keep one index have a unique meaning
**b**. Some notations lack explanations. What does the $f$ in Eq. (3) mean? What does $\hat{E}$ in Figure 1 mean? In Eq. (1), A and V should have subscripts. Also, the distribution whose mean is calculated should be listed under $E$ in Eq. (5) & (8).
**c**. Some equations have errors. In line 87-93, It would be better if the renormalization operation is taken over M rather than N. Namely, the denominator in Line should be the sum of M items since the remaining (K-M) scores are $-\infty$. In Eq. (4), "| X=X". 
**d**. For some conventional definitions like the meaning of P(A) and E(A) in line 152-153, it is not necessary to point out them in the main text.
**e**. Putting lots of formulas/notations into the paper sometimes makes the content hard to follow if there is not enough illustrations and explanations. For example, Figure 1 is hard to parse without any captions about the definitions of symbols. It could be companied with graph illustrations of attention blocks or Sparse MoE. 
- Redundant text. The preliminary and approach sections are redundant.  In Sec 1.2., there is no need to put too many words to explain the commonly-known fact that selecting top M values to calculate softmax is equivalent to the normalization over top M values after softmax. and I am not convinced that this explanation is necessary for further analyzes.  The explanations of two context-dependent expert selection methods could be condensed. The paper should be delivered in a concise way with compact information. 
- Some typos. In line 20,  "Motivated by" not "Motivating by". In line 166, it should be "Considering". In line 255, 261, 267, there are some indexes (5, 6, 7) coming from nowhere. In line 290, it should be "similar to". In 128-129 “is as” and a misplaced comma seems confusing to me.
- Others. The EMPIRICAL ANALYSIS and EXPERIMENTAL RESULTS could be merged rather than set apart.


2. The novelty of the paper is questionable. There is lack of sufficient literature review. The papers listed in the related sections seem to be "out-of-date", i.e. published two years ago and the proposed methods were compared against a limited set of long time-established baselines. Although it is interesting to see that the authors fit the MoE into PGM framework, the proposed "attention-like" aggregation techniques are still quite simple and straightforward, making me question the necessity of using this complicated PGM framework.

3. Experiments. In Figure 2, the entropy curve for the baseline is missing in the right plot, and the dynamics of fluctuation values across different layers should be analyzed.

questions:
See the above weaknesses part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732071552037

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of revision

comment:
Incorporating the insightful comments from the reviewers into our paper has helped us improve its quality. Here, we provide a summary of the revisions (highlighted in blue in the revised manuscript) for convenience:
1. We fixed typos suggested by reviewers.
2. Section 6: we added the discussion of more recent related works. 
3. Appendix D.3: we compare and combine Mutual-Inform SMoEs with previous works to confirm the benefit of our methods.
4. Appendix D.4: we evaluate the adaptability of our models, we examine the fine-tuning performance of pretrained SMoE variants on different downstream tasks.
5. Appendix D.5: we demonstrate the scalability of our models by evaluating them with a larger transformer-MoEs baseline. 
6. Appendix D.6: we evaluate performance across different model configurations, we experiment with varying numbers of experts and active experts.
7. Appendix D.7: we compare the computational complexity and memory complexity of using mutual inform techniques to the baselines.
8. Appendix C.1 and C.2: we add the model size for the baselines.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732072195620

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response (2/3)

comment:
2.**Adaptability of our models**. To evaluate the adaptability of our models, we examine the fine-tuning performance of pretrained SMoE variants on different downstream tasks. Specifically, we assess their test accuracy on Stanford Sentiment Treebank 5[1], 2[2] (SST5, SST2), and Banking-77[3] (B77) datasets, as shown in Table 2 (or see Table 4, Appendix D.4). The results demonstrate that Attention-inform SMoE consistently achieves the highest accuracy across all datasets (38.89% on SST5, 72.41% on SST2, and 85.84% on B77), while Similarity-inform SMoE also consistently outperforms the conventional SMoE baseline. These improvements in fine-tuning performance across diverse tasks suggest that our proposed Mutual-inform SMoE variants possess better adaptability compared to conventional SMoEs.

Table 2: *Top-1 test accuracy on Stanford Sentiment Treebank 5, 2 (SST5, SST2), and Banking-77 (B77) finetuning task for SMoE-medium size variants, M=2*
| Model | sst5 | sst2 | banking77 |
|------------|---------------------------|--------------------------|--------------------------------|
| SMoE | 36.54 | 70.23 | 83.96 |
| Similarity-inform SMoE | 37.91 | 71.72 | 85.19 |
| Attention-inform SMoE | **38.89** | **72.41** | **85.84** |

**References**

[1, 2]Socher et al. "*Recursive Deep Models for Semantic Compositionality Over a Sentiment Treebank.*" EMNLP. 2013.

[3]Casanueva et al. "*Efficient intent detection with dual sentence encoders.*" NLP4ConvAI. 2013.

3.**Scalability of our models** To further demonstrate the scalability of our models, we evaluate them with a larger transformer-MoEs baseline of approximately 390M parameters, with 12 layers. The results in Table 3 (or Table 5 in Appendix D.5 ) confirms that the scaling law holds true, as all models show improved language modeling performance with increased parameter count. Importantly, both Similarity-inform SMoE and Attention-inform SMoE maintain their performance advantage over the conventional SMoE at this larger scale, with Similarity-inform SMoE emerging as the best performing variant. These findings validate that the benefits of our proposed methods are preserved when scaling up model size.

Table 3: *PPL evaluation with the clean and attacked Wikitext-103 of ***SMoE-large*** size variants, M=2*
| Model/Metric | Clean Wikitext-103 Valid PPL | Clean Wikitext-103 Test PPL | Attacked Wikitext-103 Valid PPL | Attacked Wikitext-103 Test PPL |
|------------|---------------------------|--------------------------|--------------------------------|-------------------------------|
| SMoE | 28.737 | 30.378 | 36.43 | 38.34 |
| Similarity-inform SMoE | **27.06** | **28.34** | **34.65** | **36.28** |
| Attention-inform SMoE | 27.26 | 28.69 | 34.69 | 36.37 |

4.**Different number of experts and Top-M**. While our original submission included experiments with 16 experts and varying active experts (top-1, top-2), we have now added new experiments with 32 experts and top-2, as well as 16 experts and top-8 routing during the rebuttal period to provide a more comprehensive analysis of trends. Across all these settings, both Similarity-inform SMoE and Attention-inform SMoE perform consistently better than the baseline SMoE, achieving lower PPL scores on both Clean and Attacked Wikitext-103 datasets (Table 4, or Table 6 in Appendix D.6). When using 32 experts, our methods achieve PPL reductions of up to 1.56 PPL compared to the baseline, and when increasing to top-8 active experts, they maintain their advantage with improvements of up to 1.64 PPL. These performance gains further demonstrate the robustness and effectiveness of our models.

Table 4: *PPL evaluation (lower is better) with the clean and attacked Wikitext-103 of baseline SMoEs and Mutual-Inform SMoE(s) with different number of experts and Top-M.*
| Model/Metric | Clean Valid PPL | Clean Test PPL | Attacked Valid PPL | Attacked Test PPL |
|------------|---------------------------|--------------------------|--------------------------------|-------------------------------|
| SMoE (M = 1, K = 16) | 39.55 | 40.75 | 48.82 | 50.21 |
| Similarity-inform SMoE (M = 1, K = 16) | 37.78 | 39.18 | 46.93 | 48.66 |
| Attention-inform SMoE (M = 1, K = 16) | 38.02 | 39.35 | 47.20 | 48.72 |
| SMoE (M = 2, K = 16) | 33.29 | 34.84 | 41.75 | 43.59 |
| Similarity-inform SMoE (M = 2, K = 16) | **30.75** | **32.03** | **38.33** | **39.92** |
| Attention-inform SMoE (M = 2, K = 16) | 31.31 | 32.23 | 39.68 | 40.91 |
| SMoE (M = 8, K = 16) | 33.48 | 34.92 | 41.36 | 42.98 |
| Similarity-inform SMoE (M = 8, K = 16) | 32.5 | 33.81 | 40.6 | 42.37 |
| Attention-inform SMoE (M = 8, K = 16) | 31.97 | **33.28** | 39.98 | **41.45** |
| SMoE (M=2, K = 32) | 31.82 | 33.41 | **39.90** | 41.79 |
| Similarity-inform SMoE (M=2, K = 32) | 30.41 | **31.62** | 38.23 | 39.77 |
| Attention-inform SMoE (M=2, K = 32) | **30.39** | 31.85 | **37.8** | **39.65** |
</content>
</thread_object>

</thread>
<thread>
creation date:
1732072341953

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (1/2)

comment:
Dear AC and Reviewers,

Thanks for your thoughtful reviews and valuable comments, which have helped us significantly improve the paper. We are encouraged by your endorsements that: 1) Our paper is interesting and provide valuable insights (Reviewer jpxc), its strength lies in its novelty (Reviewer jpxc, Reviewer Mg8s, Reviewer azAZ), also Reviewer azAZ highlights our novelty in innovatively reinterpreting SMoE which is a fresh way of addressing the critical issue of routing challenges; (2) our work is theoretically solid (Reviewer jpxc, Reviewer Mg8s, Reviewer azAZ), with sound, rigorous and well-presented mathematical derivations and proofs backing claims (Reviewer azAZ, Reviewer Mg8s); (3) our method and frameworkare are elegant, practical, comprehensive and rigorous (Reviewer Mg8s).
Reviewer azAZ also states that “improvements in both standard and adversarial settings underscore the paper’s value in robust model deployment”. We have also updated our submission based on the reviewers' feedbacks, and we have highlighted our revision in blue.

One of the main requests from reviewers is to extend our baselines to previous works and include more diverse baseline approaches. The reviewers also raised questions about the adaptability of our Mutual-inform SMoE and its performance across different settings, such as increases in the number of experts, number of active parameters, and the overall scalability of the model. We address these questions here. 

1.**Comparisons to previous works**. To further investigate the advantages of our Mutual-inform SMoE, we compare and adapt our proposed models to X-MoE, previous work that addresses routing fluctuation in MoE models, and SMoE-dropout, another method that improves upon the standard SMoE. While both X-MoE[1] and SMoE-dropout[2] show improved performance over the standard SMoE baseline, as shown in Table 1 (or Table 3, Appendix D.3), our Mutual-inform SMoE variants still significantly outperform them, as evidenced by the lower PPL scores across both Clean and Attacked Wikitext-103 datasets. Furthermore, when we integrate our proposed methods with these models to create Similarity-inform X-MoE, Attention-inform X-MoE, Similarity-inform SMoE-dropout, and Attention-inform SMoE-dropout variants, we observe substantial improvements over their respective baselines, with consistent PPL reductions in both validation and test sets. These results demonstrate not only the superior performance of our approach but also its effectiveness as a plug-and-play solution that can enhance various MoE architectures and improve model robustness.

Table 1: *PPL evaluation (lower is better) with the clean and attacked Wikitext-103 on valid set and test set of SMoE-medium size variants, M=2*
| Model/Metric | Clean Wikitext-103 Valid PPL | Clean Wikitext-103 Test PPL | Attacked Wikitext-103 Valid PPL | Attacked Wikitext-103 Test PPL |
|------------|---------------------------|--------------------------|--------------------------------|-------------------------------|
| SMoE | 33.29 | 34.84 | 41.75 | 43.59 |
| Similarity-inform SMoE | **30.75** | **32.03** | **38.33** | **39.92** |
| Attention-inform SMoE | 31.31 | 32.23 | 39.68 | 40.91 |
| X-MoE | 33.05 | 34.49 | 41.68 | 42.96 |
| Similarity-inform X-MoE | **31.83** | **33.06** | **39.92** | **41.28** |
| Attention-inform X-MoE | 32.06 | 33.24 | 40.35 | 41.73 |
| SMoE-dropout | 33.08 | 34.67 | 41.11 | 43.09 |
| Similarity-inform SMoE-dropout | 32.47 | 33.69 | 40.6 | 41.99 |
| Attention-inform SMoE-dropout | **32.21** | **33.91** | **40.56** | **42.17** |

**References**

[1]Chi et al. "*On the Representation Collapse of Sparse Mixture of Experts*". NeurIPS. 2022

[2]Chen, et al. "*Sparse MoE as the New Dropout: Scaling Dense and Self-Slimmable Transformers*" ICLR. 2023.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a probabilistic graphical model (PGM) framework to better understand and stabilize routing decisions in Sparse Mixture of Experts (SMoE) architectures. By identifying that standard SMoE methods rely on conditional independence assumptions leading to routing fluctuations, the authors introduce Mutual-Inform SMoEs, which incorporate inter-token influences via similarity or attention. They provide proofs of reduced routing entropy and empirically show that their approach can improve robustness and reduce fluctuations on mid-scale vision and language modeling tasks. Strengths include a theoretical motivation, a novel PGM perspective on SMoEs, and initial empirical evidence that the proposed method enhances routing stability and performance. However, several weaknesses limit the paper’s impact. The experiments are conducted on relatively small-scale models and with limited training durations, making it difficult to confidently assess scalability to large-scale LLM scenarios where MoEs are most relevant. Additionally, initial versions lacked thorough baseline comparisons and crucial details, only partly remedied during the discussion phase. Given the importance of large-scale validation for MoEs would benefit significantly from more extensive experiments, deeper comparisons, and clearer documentation of hyperparameters and methodologies.