Tight Clusters Make Specialized Experts

Authors: Stefan Nielsen, Rachel Teo, Laziz Abdullaev, Tan Minh Nguyen
Primary area: foundation or frontier models, including LLMs
Keywords: Mixture of Experts, robustness, clustering

Abstract:
Sparse Mixture-of-Experts (MoE) architectures have emerged as a promising approach to decoupling model capacity from computational cost. At the core of the MoE model is the router, which learns the underlying clustering structure of the input distribution in order to send input tokens to appropriate experts. However, latent clusters may be unidentifiable in high dimension, which causes slow convergence, susceptibility to data contamination, and overall degraded representations as the router is unable to perform appropriate token-expert matching. We examine the router through the lens of clustering optimization and derive optimal feature weights that maximally identify the latent clusters. We use these weights to compute the token-expert routing assignments in an adaptively transformed space that promotes well-separated clusters, which helps identify the best-matched expert for each token. In particular, for each expert cluster, we compute a set of weights that scales features according to whether that expert clusters tightly along that feature. We term this novel router the Adaptive Clustering (AC) router. Our AC router enables the MoE model to obtain three connected benefits: 1) faster convergence, 2) better robustness to data corruption, and 3) overall performance improvement, as experts are specialized in semantically distinct regions of the input space. We empirically demonstrate the advantages of our AC router over baseline routing methods when applied on a variety of MoE backbones for language modeling and image recognition tasks in both clean and corrupted settings.

Decision: accept (poster)
Reviewer scores: 6, 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1730284721657

thread_type:
official_review

<thread_object>
creator: reviewer_2ypo

<content>
summary:
The paper presents an Adaptive Clustering (AC) router for MoE models. The Ac router offers faster convergence, better robustness, and overall performance improvement.

strengths:
The solution is theoretical based along with extensive evaluation.

weaknesses:
- The paper targets MoE experts only, some of the notations and terms seem to be used specifically in the MoE domain.

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1730380547675

thread_type:
official_review

<thread_object>
creator: reviewer_d6ym

<content>
summary:
the paper studies policies for selection of experts in sparse MoE models, and casts the problem as optimal clustering

the main contributions are 1) to derive closed form expression for the cluster weights. the formulation elegantly extends the classic topk MoE formulation using a diagonal matrix using inverse of cluster dispersion measure,  2) to theoretically bound the incorrect assignment probability and 3) prove faster convergence (though is not clear under which assumptions the prove holds

the paper presents experimental results on  text and image classification, mostly presenting results over baseline models

strengths:
- the theoretical framework is solid, with clear explanations and remarks on the main body, and abundant material (only skimmed through) in the appendix concerning proofs  of the framework 

- the paper is clearly written and motivated. 

- the experimental section tackles two complementary domains and, whereas has less critical mass than the theoretical part, it is well executed

- extensive material in the appendix complement and completes the paper

weaknesses:
# necessity of the framework

while the theoretical framework provides and the analysis in Tab 7 shows the method to only provide a rather modest additional complexity,  and the theoretical proofs shows it to be (in a sense) sufficient to solve the routing task, an open question is whether it is also necessary for the task.

otherwise stated, in the case of top-1/top-2 expert selection, an open question is whether it would be possible to significantly simplify the operations to be carried on (same way you do not need a full sort if you only need the maximum value of a vector) to attain similar performance ?

similarly, a more systematic comparison with related competitor work would better allow to appreciate the practical impact.

# structural difference of the solution 

authors do a fair job in visually shown differences in the top experts (Fig 1), but that remains at an anecdotal level.  The load balancing analysis in sec 4.3 provides some complementary view, but is too aggregated (as it looks at measures of variation, but the baseline of perfect load balancing is clearly not a target as it likely would not be attainable).

as such, assessing and quantifying the /differences/ at individual expert level, or input level would provide a factual quantitative assessment of the number of times ACMoE expert selection differs from competing selections. (ie. instead of maps in Fig, providing a distribution of the times where top-k selection differed).  

As in many cases the end-to-end performance improvement is slim (few percentage points) this intermediate assessment could allow to perceive the amount of top-k choices that were changed wrt the baselines models.


# deviation measures

it is unclear to what extent the proofs (and the practical performance) of the framework depends on subtle hyperparameter choices that are fixed and  not considered in the paper: e.g., measure of dispersion used in sec 3 is the mean absolute deviation, which is easier to compute but less robust than median absolute deviation (admittedly, this would go in the opposite direction as for simplification; but is unclear to what extent practical details as this one can have a performance impact in practice).

to have no impact on the selection, the top-k rank would need to be maintained across a family of measures of dispersion -- the impact of which is far from being clear after reading the paper. If the proof holds for families with given properties, the paper could be reinforced by making that clear.

questions:
no further questions
</content>
</thread_object>

</thread>
<thread>
creation date:
1730644468032

thread_type:
official_review

<thread_object>
creator: reviewer_ldye

<content>
summary:
This paper proposes a novel routing mechanism called Adaptive Clustering (AC) router for Mixture-of-Experts (MoE) architectures. The key idea is to compute token-expert assignments in an adaptively transformed space that better reveals latent clusters in the data. The transformation is derived from a feature-weighted clustering optimization perspective, where features that promote tight clustering for each expert are upweighted. The authors demonstrate both theoretical and empirical advantages of their method, showing improved convergence speed, robustness to data contamination, and overall performance across language modeling and image classification tasks.

strengths:
1. The paper introduces a novel perspective by viewing the MoE routing mechanism through the lens of feature-weighted clustering optimization. This theoretical framing allows for the derivation of optimal feature weights that improve cluster separability.

2. The authors provide rigorous theoretical analysis that support their claims on the improved robustness and convergence speed of their method.

3. Extensive experiments on large-scale datasets for both language modeling and image classification demonstrate the effectiveness of the AC router.

4. The proposed method requires no additional learnable parameters and introduces negligible computational overhead. And it can be integrated into any existing MoE architecture

weaknesses:
1. The AC router relies on expert assignments from the previous layer to compute the adaptive transformation. This dependence may limit its applicability in scenarios where embedding sizes change between layers or in networks without consistent layer structures.

2. The effectiveness of the AC router is closely tied to the quality of expert assignments in earlier layers. Since the method relies on routing decisions from prior layers, it may be less effective if the early layers have not yet learned meaningful or well-separated cluster structures. In practice, early layers in deep networks often learn basic features that may not be semantically distinct enough to form tight clusters. As a result, the AC router may struggle to make optimal routing decisions in later layers if the cluster assignments in the initial layers are poor.

3. The method assumes that the input data naturally clusters in the feature space. While this assumption holds in many cases (e.g., language modeling and image classification), it might not generalize to tasks where the input data is more uniformly distributed or lacks clear clustering patterns.

questions:
1. The choice of mean absolute deviation (MAD) as the measure of dispersion is justified, but the paper does not deeply explore the impact of other potential measures, such as variance or interquartile range, on performance. Have you considered or tested other measures of dispersion besides MAD? How sensitive is the method to the choice of dispersion measure, and could alternative measures potentially improve performance or robustness?

2. How does the AC router handle situations where the assignments from previous layers are noisy or not well-defined? Is there a way to initialize or adjust the method to be effective in early layers or when prior assignments are unreliable?

3. Since the method requires consistent embedding dimensions between layers, how can it be adapted for architectures where the embedding size changes, such as in some convolutional neural networks or transformers with variable dimensions?

4. The theoretical analysis relies on Gaussian mixture models. How does the AC router perform when the data clusters have non-Gaussian distributions or are not well-separated?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730899350057

thread_type:
official_review

<thread_object>
creator: reviewer_uqyp

<content>
summary:
This paper proposes an Adaptive Clustering (AC) router for Sparse Mixture-of-Experts (MoE) architectures to improve routing efficiency. It focuses on optimizing the matching of tokens to experts by adaptively transforming the input feature space. This method promotes faster convergence, increased robustness against data contamination, and improved overall performance. The authors demonstrate the AC router's effectiveness in large-scale language and image tasks, outperforming baseline routers in robustness and efficiency without added learnable parameters.

strengths:
1.	The proposed router uses a unique, feature-weighted approach for adaptive clustering, enabling more efficient specialization in MoE models. It achieves improvements without extra learnable parameters or added computation cost, making it ideal for scalable applications.

2.	The paper provides robust theoretical proof supporting the method's faster convergence and increased robustness

3.	The experiments cover various large-scale tasks, such as WikiText-103 and ImageNet, and include backbones like Switch Transformer, GLaM, and Swin Transformers, evaluating performance under both clean and corrupted data conditions.

weaknesses:
1.	The approach seems to tailor to settings with well-structured clusters, which may limit performance on datasets or tasks where latent clusters are less distinct or not aligned with chosen features.

2.	The AC router relies on prior layers for initial expert assignments, potentially constraining flexibility and requiring consistent embedding sizes across layers.

3.	Although tested on large-scale datasets, further application to more dynamic datasets could strengthen claims of robustness and adaptability.

questions:
1.	How does the AC router handle scenarios with overlapping or dynamically evolving clusters where expert specialization may be less clear?

2.	What mechanisms could be integrated to enhance adaptability in routing without relying on previous layer assignments?
</content>
</thread_object>

</thread>
<thread>
creation date:
1731942329475

thread_type:
official_comment

<thread_object>
creator: authors

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

comment:
Dear AC and reviewers,

Thanks for your thoughtful reviews and valuable comments, which have helped us improve the paper significantly. We are encouraged by the endorsements that: 1) our proposed framework and method are unique and novel (reviewers UqYP, LDye), 2) our theoretical analysis is rigorous (reviewers UqYP, LDye, D6YM), 3) our experimental evaluation is comprehensive (all reviewers), and 4) our method's efficiency makes it easily scalable and integrable into any MoE architecture (reviewers UqYP, LDye).

One common question emerging from reviewers was regarding practical limitations of our method given it requires the embedding size to remain constant across adjacent MoE layers.  Another shared question was on the adaptibility of our method to situations in which latent clusters are not well-defined or the estimated cluster assignments coming from previous layers are noisy. We address these questions here.

**Common embedding size design choices in contemporary MoE architectures.** In the majority of contemporary deep MoE architectures, the embedding size is typically constant throughout the entire model, and examples of changing embedding sizes are more the exception than the rule. For example, in Switch [3], GLaM [4], SoftMoE [1], and VMoE [5], the embedding size remains the same throughout. Swin [2] is, so far, the only transformer-MoE model we've encountered that features a changing embedding size, but still maintains the same embedding size for 18 of its 22 total layers, and so there remains ample opportunity for applying ACMoE. 

We note further that many standard-practice architectural features require the same assumption of constant embedding size, such as the residual connection. So in this sense, our requirement of constant embedding size at adjacent MoE layers is no more restrictive than what is required for commonplace design choices.

**ACMoE's adaptibility to settings with overlapping clusters and/or noisy estimated cluster assignments.**
The setting of unreliable or noisy estimated cluster assignments coming from previous layers, and the setting of overlapping clusters can be handled through a straightforward extension of our framework in which we mix the cluster-wise feature weights with mixing proportions corresponding to the affinities in the routing. For example, in a top-2 setting, if $\boldsymbol h$ has affinity scores $\alpha$ and $1-\alpha$ to clusters $k$ and $k'$ respectively, then we could also obtain the required AC routing transformation (Definition 1 in Section 3.1) for $\boldsymbol h$ as $\boldsymbol M_{k^*} = \alpha \boldsymbol M_{k} + (1-\alpha)\boldsymbol M_{k'}$. This approach therefore factors in the confidence with which we believe $\boldsymbol h$ belongs to cluster $k$ or $k'$, and so can be used in settings where we are less sure about the cluster assignment of $\boldsymbol h$ and would prefer to factor in this uncertainty. Furthermore, this adaptation of ACMoE naturally accomodates higher expert granularity backbones (i.e higher top-k settings) or SoftMoE, where all experts are active for every token. In this setting, we have many, overlapping expert clusters and we prefer our routing transformation to model the fact that each token may originate from numerous clusters.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731942481402

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Revisions

comment:
Incorporating comments and suggestions from reviewers, as well as some further empirical studies we believe informative, we summarize here the main changes in the revised paper:

1. We have conducted additional experiments on ACMoE with cluster weight mixing (Appendix C.6). We show in Tables 15 and 16 of Appendix C.6 the results of weight mixing over the top-2 highest affinity experts in Switch and GLaM backbones, where we see similar performance in Switch and a large improvement in GLaM. This straightforward extension of our framework can be used to factor in the confidence with which we believe a token belongs to an expert cluster, and therefore is useful if we believe previous layer expert assignments are noisy. Furthermore, this setup is useful for integrating ACMoE into higher granularity backbones (where we wish to activate a larger number of experts per token), such as SoftMoE, which we discuss in the next point.
2. We have conducted additional experiments on ACMoE in the SoftMoE backbone (Appendix C.7). We present in Tables 17 and 18 of Appendix C.7 performance gains of ACMoE over SoftMoE on clean, adversarially attacked, and out-of-distribution ImageNet-1K, where ACMoE delivers substantial robust performance improvements in the range of 6-7%.
3. We have conducted additional experiments on Swin-ACMoE in a larger 0.5B parameter 'Base' configuration (Appendix C.8). Table 19 in Appendix C.8 shows ACMoE continues to deliver consistent gains in the larger configuration, with robust performance improvements in the range of 3%. 
4. We have conducted an empirical assessment of the routing stability (proportion of tokens for which the expert assignments change as the tokens pass through the model) of SMoE, XMoE, StableMoE, and ACMoE on large-scale language modeling in the Switch Transformer backbone. The assessment and details can be found in Figure 5 in Appendix C.9. We see that for a trained model, ACMoE is substantially better at maintaining consistent routing through the model.
5. We have conducted additional baseline experiments using StableMoE [1] at the Switch and GLaM medium configuration in order to add further empirical support for the proposed benefits of our AC routing scheme. We add these results to Table 3 on page 8.
6. We have conducted an additional ablation study in Tables 13 and 14 in Appendix C.5.3 where we replace the diagonal elements of the AC routing transformation with mean 1 normal random variables. Though unrequested, we nonetheless thought such a study offers useful empirical insight and may be of interest to reviewers, as it shows that the performance gains brought about by ACMoE are not simply the result of noise-induced regularization. Tested in large-scale language modeling and image classification, we see in both Swin and Switch that ACMoE substantially outperforms the random ablation model.
7. We have conducted additional experiments in dynamic mixture of experts [2] where we integrate ACMoE into dynamic top-p routing in the Switch backbone. Results can be found in Table 20 of Appendix C.10.

[1] Dai et al. StableMoE: Stable Routing Strategy for Mixture of Experts (ACL 2022)

[2] Guo et al. Dynamic mixture of experts: An auto-tuning approach for efficient transformer models (2024)
</content>
</thread_object>

</thread>
<thread>
creation date:
1732534948481

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Additional Experimental Results on Dynamic Mixture of Experts

comment:
Based on reviewer UqYP's suggestion, we have further tested our AC routing framework in dynamic MoE settings. 

Following the approach of [1], we implement ACMoE into top-$p$ dynamic gating. In this setting, rather than routing each token to its top-k highest affinity experts in each MoE layer, we route each token to all experts that have affinity over a certain threshold $p$. This setting permits dynamically activating varying numbers of experts for different tokens at different layers throughout the model. We integrate our AC routing directly into this setting using the same setup as in Section 3 of our manuscript, where the AC routing transformation is computed based on the estimated cluster membership of each token using the top affinity assignment of the previous layer. We present the results for Switch transformer on WikiText-103 language modeling in the following Table A. The same results can be found in Table 20 in Appendix C.10 of our revised manuscript.

**Table A: Results on Top-p Dynamic Routing in Switch Backbone**  

| **Model**                        | **Test PPL (↓)** |
|----------------------------------|------------------|
| **Fixed top-k routing [2]**          |                  |
| SMoE-medium (Shazeer et al., 2017) | 35.48           |
| ACMoE-medium (**Ours**)          | **34.42**        |
|                                  |                  |
| **Dynamic top-$p$ routing [1]**        |                  |
| Switch-Fixed $p$                   | 35.20           |
| Switch-ACMoE-Fixed $p$ (**Ours**)  | **34.14**        |
| Switch-Learnable $p$               | 34.29           |
| Switch-ACMoE-Learnable $p$ (**Ours**) | **33.49**     |

For fixed $p$, we set $p = 0.05$. For learnable $p$, we initialize the parameter to 0.05. We select this initialization as it reproduces approximately similar performance in the Switch backbone under default top-2 routing, thereby aiding direct comparison between fixed top-k and dynamic top-$p$ routing. We see in the dynamic routing setting, ACMoE maintains the same consistent improvement over the Switch baseline of roughly 1 full PPL. These results suggest ACMoE is well-suited to the dynamic MoE setting.

We have correspondingly updated our summary of revisions and the current uploaded manuscript.

[1] Guo et al. Dynamic mixture of experts: An auto-tuning approach for efficient transformer models (2024)

[2] Shazeer et al. The Sparsely-Gated Mixture-of-Experts Layer (ICLR 2017)
</content>
</thread_object>

</thread>
<thread>
creation date:
1732739395381

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Our Discussion with the Reviewers So Far

comment:
Dear Reviewers, Senior Area Chairs, and Area Chairs,

We would like to summarize the revisions we have made so far, incorporating additional results and improvements based on the reviewers' suggestions:

**- [Reviewers UqYP, LDye, D6YM]** We have conducted additional experiments using ACMoE in the SoftMoE backbone evaluated on ImageNet. We present in Tables 17 and 18 of Appendix C.7 performance gains of ACMoE over SoftMoE on clean, adversarially attacked, and out-of-distribution data, where ACMoE delivers substantial robust performance improvements in the range of 6-7%. Incorporating ACMoE into SoftMoE addressed the question of how adaptable our AC routing framework is to settings where latent clusters are modeled as overlapping or indistinct, and adds further justification to the compatibility of our framework with a wide range of MoE backbones.

**- [Reviewers UqYP, LDye]** We have conducted additional experiments using ACMoE in GlaM, Switch, and Swin backbones using cluster weight mixing, which is a straightforward extension of our framework to handle situations where previous layer cluster assignments may be noisy or unreliable. We present the results in Tables 15 and 16 of Appendix C.6, where we see that ACMoE maintains strong performance or even improves with an additional whole PPL in the case of GLaM.

**- [Reviewer UqYP]** We have conducted additional experiments using ACMoE in the Switch backbone using dynamic top-$p$ routing in order to empirically justify the adaptability of ACMoE to the dynamic MoE setting. We present the results in Table 20 in Appendix C.10, where we see ACMoE maintains the same strong, consistent performance gains over Switch transformer.

**- [Reviewer D6YM]** We have conducted additional baseline experiments using StableMoE in the Switch and GLaM medium configuration in order to add further empirical support for the proposed benefits of our AC routing scheme over baseline methods. Results are found in Table 3 of the main text.

**- [Reviewer D6YM]** We have conducted an empirical assessment of the routing stability (proportion of tokens for which the expert assignments change as the tokens pass through the model) of SMoE, XMoE, StableMoE, and ACMoE in the Switch Transformer backbone. The assessment and details can be found in Figure 5 in Appendix C.9. We see that for a trained model, ACMoE is substantially better at maintaining consistent routing through the model.

**- [Reviewers D6YM, LDye]** We have provided ablation studies on layer-wise placement of ACMoE (Tables 10 & 12, Appendix C.5.2), which show that the AC router is able to improve token-expert matching even at early layers in the network. We also ablate the measure of dispersion used in the AC routing transformation (Tables 9 & 11, Appendix C.5.1), finding that the framework is robust to the selected measure of dispersion, but attains top performance when using MAD as reported.

**- [Reviewer LDye]** We have provided theoretical clarification for the justification and broad applicability of our GMM modeling setup used in our robustness propositions.

**- [Reviewers LDye, D6YM, UqYP]** We have provided additional justification for the widespread and easy incorporability of our method into contemporary MoE architectures, and also further justified the benefits of estimating the tokens' cluster membership by previous layer assignments by appealing to the considerable efficiency advantages of the scheme.

**- [Provided without request]** To further enhance the empirical justification of our method, we additionally provided clean and robust results for ACmoE in the larger 0.5B Swin backbone to demonstrate the ability of our method to scale (Table 19, Appendix C.8). We also conducted a random ablation study to demonstrate the performance gains of ACMoE do not stem from noise-induced regularization (Tables 13 & 14, Appendix C.5.3).

We thank the reviewers for their valuable comments, which have helped us improve our paper both theoretically and empirically. We are happy to address any additional comments or questions during the extended discussion phase.

Best regards,

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces the Adaptive Clustering (AC) router for Mixture-of-Experts (MoE), providing a clustering-based approach to token-expert routing. The method achieves faster convergence, improved robustness, and better performance without added parameters. The theoretical analysis is rigorous, and extensive experiments on language and vision tasks demonstrate significant gains over baselines. The paper is well-written, technically sound, and makes a valuable contribution to MoE research.